{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "524620c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp core"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15392f6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "12fa25a4",
   "metadata": {},
   "source": [
    "# Core\n",
    "> NeuralForecast contains two main components, PyTorch implementations deep learning predictive models, as well as parallelization and distributed computation utilities. The first component comprises low-level PyTorch model estimator classes like `models.NBEATS` and `models.RNN`. The second component is a high-level `core.NeuralForecast` wrapper class that operates with sets of time series data stored in pandas DataFrames."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "515672ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "import os\n",
    "import shutil\n",
    "import sys\n",
    "\n",
    "import git\n",
    "import s3fs\n",
    "from fastcore.test import test_eq, test_fail\n",
    "from nbdev.showdoc import show_doc\n",
    "from neuralforecast.utils import generate_series\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44065066-e72a-431f-938f-1528adef9fe8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "import pickle\n",
    "import warnings\n",
    "from copy import deepcopy\n",
    "from itertools import chain\n",
    "from typing import Any, Dict, List, Optional, Sequence, Union\n",
    "\n",
    "import fsspec\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pytorch_lightning as pl\n",
    "import torch\n",
    "import utilsforecast.processing as ufp\n",
    "from coreforecast.grouped_array import GroupedArray\n",
    "from coreforecast.scalers import (\n",
    "    LocalBoxCoxScaler,\n",
    "    LocalMinMaxScaler,\n",
    "    LocalRobustScaler,\n",
    "    LocalStandardScaler,\n",
    ")\n",
    "from utilsforecast.compat import DataFrame, DFType, Series, pl_DataFrame, pl_Series\n",
    "from utilsforecast.validation import validate_freq\n",
    "\n",
    "from neuralforecast.common._base_model import DistributedConfig\n",
    "from neuralforecast.compat import SparkDataFrame\n",
    "from neuralforecast.losses.pytorch import HuberIQLoss, IQLoss\n",
    "from neuralforecast.tsdataset import _FilesDataset, TimeSeriesDataset, LocalFilesTimeSeriesDataset\n",
    "from neuralforecast.models import (\n",
    "    GRU, LSTM, RNN, TCN, DeepAR, DilatedRNN,\n",
    "    MLP, NHITS, NBEATS, NBEATSx, DLinear, NLinear,\n",
    "    TFT, VanillaTransformer,\n",
    "    Informer, Autoformer, FEDformer,\n",
    "    StemGNN, PatchTST, TimesNet, TimeLLM, TSMixer, TSMixerx,\n",
    "    MLPMultivariate, iTransformer,\n",
    "    BiTCN, TiDE, DeepNPTS, SOFTS,\n",
    "    TimeMixer, KAN, RMoK, TimeXer\n",
    ")\n",
    "from neuralforecast.common._base_auto import BaseAuto, MockTrial\n",
    "from neuralforecast.utils import PredictionIntervals, get_prediction_interval_method, level_to_quantiles, quantiles_to_level"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb8b4b3c-04bf-4a92-9a1a-b60735997c36",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "# this disables warnings about the number of workers in the dataloaders\n",
    "# which the user can't control\n",
    "warnings.filterwarnings(\"ignore\", category=pl.utilities.warnings.PossibleUserWarning)\n",
    "\n",
    "def _insample_times(\n",
    "    times: np.ndarray,\n",
    "    uids: Series,\n",
    "    indptr: np.ndarray,\n",
    "    h: int,\n",
    "    freq: Union[int, str, pd.offsets.BaseOffset],\n",
    "    step_size: int = 1,\n",
    "    id_col: str = 'unique_id',\n",
    "    time_col: str = 'ds',\n",
    ") -> DataFrame:\n",
    "    sizes = np.diff(indptr)\n",
    "    if (sizes < h).any():\n",
    "        raise ValueError('`sizes` should be greater or equal to `h`.')\n",
    "    # TODO: we can just truncate here instead of raising an error\n",
    "    ns, resids = np.divmod(sizes - h, step_size)\n",
    "    if (resids != 0).any():\n",
    "        raise ValueError('`sizes - h` should be multiples of `step_size`')\n",
    "    windows_per_serie = ns + 1\n",
    "    # determine the offsets for the cutoffs, e.g. 2 means the 3rd training date is a cutoff\n",
    "    cutoffs_offsets = step_size * np.hstack([np.arange(w) for w in windows_per_serie])\n",
    "    # start index of each serie, e.g. [0, 17] means the the second serie starts on the 18th entry\n",
    "    # we repeat each of these as many times as we have windows, e.g. windows_per_serie = [2, 3]\n",
    "    # would yield [0, 0, 17, 17, 17]\n",
    "    start_idxs = np.repeat(indptr[:-1], windows_per_serie)\n",
    "    # determine the actual indices of the cutoffs, we repeat the cutoff for the complete horizon\n",
    "    # e.g. if we have two series and h=2 this could be [0, 0, 1, 1, 17, 17, 18, 18]\n",
    "    # which would have the first two training dates from each serie as the cutoffs\n",
    "    cutoff_idxs = np.repeat(start_idxs + cutoffs_offsets, h)\n",
    "    cutoffs = times[cutoff_idxs]\n",
    "    total_windows = windows_per_serie.sum()\n",
    "    # determine the offsets for the actual dates. this is going to be [0, ..., h] repeated\n",
    "    ds_offsets = np.tile(np.arange(h), total_windows)\n",
    "    # determine the actual indices of the times\n",
    "    # e.g. if we have two series and h=2 this could be [0, 1, 1, 2, 17, 18, 18, 19]\n",
    "    ds_idxs = cutoff_idxs + ds_offsets\n",
    "    ds = times[ds_idxs]\n",
    "    if isinstance(uids, pl_Series):\n",
    "        df_constructor = pl_DataFrame\n",
    "    else:\n",
    "        df_constructor = pd.DataFrame\n",
    "    out = df_constructor(\n",
    "        {\n",
    "            id_col: ufp.repeat(uids, h * windows_per_serie),\n",
    "            time_col: ds,\n",
    "            'cutoff': cutoffs,\n",
    "        }\n",
    "    )\n",
    "    # the first cutoff is before the first train date\n",
    "    actual_cutoffs = ufp.offset_times(out['cutoff'], freq, -1)\n",
    "    out = ufp.assign_columns(out, 'cutoff', actual_cutoffs)\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6aead6db-170a-4a74-baf3-7cbaf8b7468c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "uids = pd.Series(['id_0', 'id_1'])\n",
    "indptr = np.array([0, 4, 10], dtype=np.int32)\n",
    "h = 2\n",
    "for step_size, freq, days in zip([1, 2], ['D', 'W-THU'], [1, 14]):\n",
    "    times = np.hstack([\n",
    "        pd.date_range('2000-01-01', freq=freq, periods=4),\n",
    "        pd.date_range('2000-10-10', freq=freq, periods=10),\n",
    "    ])    \n",
    "    times_df = _insample_times(times, uids, indptr, h, freq, step_size=step_size)\n",
    "    pd.testing.assert_frame_equal(\n",
    "        times_df.groupby('unique_id')['ds'].min().reset_index(),\n",
    "        pd.DataFrame({\n",
    "            'unique_id': uids,\n",
    "            'ds': times[indptr[:-1]],\n",
    "        })\n",
    "    )\n",
    "    pd.testing.assert_frame_equal(\n",
    "        times_df.groupby('unique_id')['ds'].max().reset_index(),\n",
    "        pd.DataFrame({\n",
    "            'unique_id': uids,\n",
    "            'ds': times[indptr[1:] - 1],\n",
    "        })\n",
    "    )\n",
    "    cutoff_deltas = (\n",
    "        times_df\n",
    "        .drop_duplicates(['unique_id', 'cutoff'])\n",
    "        .groupby('unique_id')\n",
    "        ['cutoff']\n",
    "        .diff()\n",
    "        .dropna()\n",
    "    )\n",
    "    assert cutoff_deltas.nunique() == 1\n",
    "    assert cutoff_deltas.unique()[0] == pd.Timedelta(f'{days}D')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c58a8a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "MODEL_FILENAME_DICT = {\n",
    "    'autoformer': Autoformer, 'autoautoformer': Autoformer,\n",
    "    'deepar': DeepAR, 'autodeepar': DeepAR,\n",
    "    'dlinear': DLinear, 'autodlinear': DLinear,\n",
    "    'nlinear': NLinear, 'autonlinear': NLinear,    \n",
    "    'dilatedrnn': DilatedRNN , 'autodilatedrnn': DilatedRNN,\n",
    "    'fedformer': FEDformer, 'autofedformer': FEDformer,\n",
    "    'gru': GRU, 'autogru': GRU,\n",
    "    'informer': Informer, 'autoinformer': Informer,\n",
    "    'lstm': LSTM, 'autolstm': LSTM,\n",
    "    'mlp': MLP, 'automlp': MLP,\n",
    "    'nbeats': NBEATS, 'autonbeats': NBEATS,\n",
    "    'nbeatsx': NBEATSx, 'autonbeatsx': NBEATSx,\n",
    "    'nhits': NHITS, 'autonhits': NHITS,\n",
    "    'patchtst': PatchTST, 'autopatchtst': PatchTST,\n",
    "    'rnn': RNN, 'autornn': RNN,\n",
    "    'stemgnn': StemGNN, 'autostemgnn': StemGNN,\n",
    "    'tcn': TCN, 'autotcn': TCN, \n",
    "    'tft': TFT, 'autotft': TFT,\n",
    "    'timesnet': TimesNet, 'autotimesnet': TimesNet,\n",
    "    'vanillatransformer': VanillaTransformer, 'autovanillatransformer': VanillaTransformer,\n",
    "    'timellm': TimeLLM,\n",
    "    'tsmixer': TSMixer, 'autotsmixer': TSMixer,\n",
    "    'tsmixerx': TSMixerx, 'autotsmixerx': TSMixerx,\n",
    "    'mlpmultivariate': MLPMultivariate, 'automlpmultivariate': MLPMultivariate,\n",
    "    'itransformer': iTransformer, 'autoitransformer': iTransformer,\n",
    "    'bitcn': BiTCN, 'autobitcn': BiTCN,\n",
    "    'tide': TiDE, 'autotide': TiDE,\n",
    "    'deepnpts': DeepNPTS, 'autodeepnpts': DeepNPTS,\n",
    "    'softs': SOFTS, 'autosofts': SOFTS,\n",
    "    'timemixer': TimeMixer, 'autotimemixer': TimeMixer,\n",
    "    'kan': KAN, 'autokan': KAN,\n",
    "    'rmok': RMoK, 'autormok': RMoK,\n",
    "    'timexer': TimeXer, 'autotimexer': TimeXer\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c621a39-5658-4850-95c4-050eee97403d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "_type2scaler = {\n",
    "    'standard': LocalStandardScaler,\n",
    "    'robust': lambda: LocalRobustScaler(scale='mad'),\n",
    "    'robust-iqr': lambda: LocalRobustScaler(scale='iqr'),\n",
    "    'minmax': LocalMinMaxScaler,\n",
    "    'boxcox': lambda: LocalBoxCoxScaler(method='loglik', lower=0.0)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4dae43c-4d11-4bbc-a431-ac33b004859a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class NeuralForecast:\n",
    "    \n",
    "    def __init__(\n",
    "        self, \n",
    "        models: List[Any],\n",
    "        freq: Union[str, int],\n",
    "        local_scaler_type: Optional[str] = None\n",
    "    ):\n",
    "        \"\"\"\n",
    "        The `core.StatsForecast` class allows you to efficiently fit multiple `NeuralForecast` models \n",
    "        for large sets of time series. It operates with pandas DataFrame `df` that identifies series \n",
    "        and datestamps with the `unique_id` and `ds` columns. The `y` column denotes the target \n",
    "        time series variable.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        models : List[typing.Any]\n",
    "            Instantiated `neuralforecast.models` \n",
    "            see [collection here](https://nixtla.github.io/neuralforecast/models.html).\n",
    "        freq : str or int\n",
    "            Frequency of the data. Must be a valid pandas or polars offset alias, or an integer.\n",
    "        local_scaler_type : str, optional (default=None)\n",
    "            Scaler to apply per-serie to all features before fitting, which is inverted after predicting.\n",
    "            Can be 'standard', 'robust', 'robust-iqr', 'minmax' or 'boxcox'\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        self : NeuralForecast\n",
    "            Returns instantiated `NeuralForecast` class.\n",
    "        \"\"\"\n",
    "        assert all(model.h == models[0].h for model in models), 'All models should have the same horizon'\n",
    "\n",
    "        self.h = models[0].h\n",
    "        self.models_init = models\n",
    "        self.freq = freq\n",
    "        if local_scaler_type is not None and local_scaler_type not in _type2scaler:\n",
    "            raise ValueError(f'scaler_type must be one of {_type2scaler.keys()}')\n",
    "        self.local_scaler_type = local_scaler_type\n",
    "        self.scalers_: Dict\n",
    "\n",
    "        # Flags and attributes\n",
    "        self._fitted = False\n",
    "        self._reset_models()\n",
    "        self._add_level = False\n",
    "\n",
    "    def _scalers_fit_transform(self, dataset: TimeSeriesDataset) -> None:\n",
    "        self.scalers_ = {}        \n",
    "        if self.local_scaler_type is None:\n",
    "            return None\n",
    "        for i, col in enumerate(dataset.temporal_cols):\n",
    "            if col == 'available_mask':\n",
    "                continue\n",
    "            ga = GroupedArray(dataset.temporal[:, i].numpy(), dataset.indptr)                \n",
    "            self.scalers_[col] = _type2scaler[self.local_scaler_type]().fit(ga)\n",
    "            dataset.temporal[:, i] = torch.from_numpy(self.scalers_[col].transform(ga))\n",
    "\n",
    "    def _scalers_transform(self, dataset: TimeSeriesDataset) -> None:\n",
    "        if not self.scalers_:\n",
    "            return None\n",
    "        for i, col in enumerate(dataset.temporal_cols):\n",
    "            scaler = self.scalers_.get(col, None)\n",
    "            if scaler is None:\n",
    "                continue\n",
    "            ga = GroupedArray(dataset.temporal[:, i].numpy(), dataset.indptr)\n",
    "            dataset.temporal[:, i] = torch.from_numpy(scaler.transform(ga))\n",
    "\n",
    "    def _scalers_target_inverse_transform(self, data: np.ndarray, indptr: np.ndarray) -> np.ndarray:\n",
    "        if not self.scalers_:\n",
    "            return data\n",
    "        for i in range(data.shape[1]):\n",
    "            ga = GroupedArray(data[:, i], indptr)\n",
    "            data[:, i] = self.scalers_[self.target_col].inverse_transform(ga)\n",
    "        return data\n",
    "\n",
    "    def _prepare_fit(self, df, static_df, predict_only, id_col, time_col, target_col):\n",
    "        #TODO: uids, last_dates and ds should be properties of the dataset class. See github issue.\n",
    "        self.id_col = id_col\n",
    "        self.time_col = time_col\n",
    "        self.target_col = target_col\n",
    "        self._check_nan(df, static_df, id_col, time_col, target_col)\n",
    "        \n",
    "        dataset, uids, last_dates, ds = TimeSeriesDataset.from_df(\n",
    "            df=df,\n",
    "            static_df=static_df,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "            target_col=target_col,\n",
    "        )\n",
    "        if predict_only:\n",
    "            self._scalers_transform(dataset)\n",
    "        else:\n",
    "            self._scalers_fit_transform(dataset)\n",
    "        return dataset, uids, last_dates, ds\n",
    "\n",
    "\n",
    "    def _check_nan(self, df, static_df, id_col, time_col, target_col):\n",
    "        cols_with_nans = []\n",
    "\n",
    "        temporal_cols = [target_col] + [c for c in df.columns if c not in (id_col, time_col, target_col)]\n",
    "        if \"available_mask\" in temporal_cols:\n",
    "            available_mask = df[\"available_mask\"].to_numpy().astype(bool)\n",
    "        else:\n",
    "            available_mask = np.full(df.shape[0], True)\n",
    "\n",
    "        df_to_check = ufp.filter_with_mask(df, available_mask)\n",
    "        for col in temporal_cols:\n",
    "            if ufp.is_nan_or_none(df_to_check[col]).any():\n",
    "                cols_with_nans.append(col)\n",
    "\n",
    "        if static_df is not None:\n",
    "            for col in [x for x in static_df.columns if x != id_col]:\n",
    "                if ufp.is_nan_or_none(static_df[col]).any():\n",
    "                    cols_with_nans.append(col)\n",
    "\n",
    "        if cols_with_nans:\n",
    "            raise ValueError(f\"Found missing values in {cols_with_nans}.\")\n",
    "\n",
    "    def _prepare_fit_distributed(\n",
    "        self,\n",
    "        df: SparkDataFrame,\n",
    "        static_df: Optional[SparkDataFrame],\n",
    "        id_col: str,\n",
    "        time_col: str,\n",
    "        target_col: str,\n",
    "        distributed_config: Optional[DistributedConfig],\n",
    "    ):\n",
    "        if distributed_config is None:\n",
    "            raise ValueError(\n",
    "                \"Must set `distributed_config` when using a spark dataframe\"\n",
    "            )\n",
    "        if self.local_scaler_type is not None:\n",
    "            raise ValueError(\n",
    "                \"Historic scaling isn't supported in distributed. \"\n",
    "                \"Please open an issue if this would be valuable to you.\"\n",
    "            )\n",
    "        temporal_cols = [c for c in df.columns if c not in (id_col, time_col)]\n",
    "        if static_df is not None:\n",
    "            static_cols = [c for c in static_df.columns if c != id_col]\n",
    "            df = df.join(static_df, on=[id_col], how=\"left\")\n",
    "        else:\n",
    "            static_cols = None\n",
    "        self.id_col = id_col\n",
    "        self.time_col = time_col\n",
    "        self.target_col = target_col\n",
    "        self.scalers_ = {}\n",
    "        num_partitions = distributed_config.num_nodes * distributed_config.devices\n",
    "        df = df.repartitionByRange(num_partitions, id_col)\n",
    "        df.write.parquet(path=distributed_config.partitions_path, mode=\"overwrite\")\n",
    "        fs, _, _ = fsspec.get_fs_token_paths(distributed_config.partitions_path)\n",
    "        protocol = fs.protocol \n",
    "        if isinstance(protocol, tuple):\n",
    "            protocol = protocol[0]\n",
    "        files = [\n",
    "            f'{protocol}://{file}'\n",
    "            for file in fs.ls(distributed_config.partitions_path)\n",
    "            if file.endswith(\"parquet\")\n",
    "        ]\n",
    "        return _FilesDataset(\n",
    "            files=files,\n",
    "            temporal_cols=temporal_cols,\n",
    "            static_cols=static_cols,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "            target_col=target_col,\n",
    "            min_size=df.groupBy(id_col).count().agg({\"count\": \"min\"}).first()[0],\n",
    "        )\n",
    "    \n",
    "    def _prepare_fit_for_local_files(\n",
    "        self, \n",
    "        files_list: Sequence[str], \n",
    "        static_df: Optional[DataFrame], \n",
    "        id_col: str, \n",
    "        time_col: str, \n",
    "        target_col: str\n",
    "    ):\n",
    "        if self.local_scaler_type is not None:\n",
    "            raise ValueError(\n",
    "                \"Historic scaling isn't supported when the dataset is split between files. \"\n",
    "                \"Please open an issue if this would be valuable to you.\"\n",
    "            )\n",
    "        \n",
    "        self.id_col = id_col\n",
    "        self.time_col = time_col\n",
    "        self.target_col = target_col   \n",
    "        self.scalers_ = {}     \n",
    "\n",
    "        exogs = self._get_needed_exog() \n",
    "        return LocalFilesTimeSeriesDataset.from_data_directories(\n",
    "            directories=files_list,\n",
    "            static_df=static_df,\n",
    "            exogs=exogs,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "            target_col=target_col,\n",
    "        )\n",
    "\n",
    "\n",
    "    def fit(\n",
    "        self,\n",
    "        df: Optional[Union[DataFrame, SparkDataFrame, Sequence[str]]] = None,\n",
    "        static_df: Optional[Union[DataFrame, SparkDataFrame]] = None,\n",
    "        val_size: Optional[int] = 0,\n",
    "        use_init_models: bool = False,\n",
    "        verbose: bool = False,\n",
    "        id_col: str = 'unique_id',\n",
    "        time_col: str = 'ds',\n",
    "        target_col: str = 'y',\n",
    "        distributed_config: Optional[DistributedConfig] = None,\n",
    "        prediction_intervals: Optional[PredictionIntervals] = None,\n",
    "    ) -> None:\n",
    "        \"\"\"Fit the core.NeuralForecast.\n",
    "\n",
    "        Fit `models` to a large set of time series from DataFrame `df`.\n",
    "        and store fitted models for later inspection.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        df : pandas, polars or spark DataFrame, or a list of parquet files containing the series, optional (default=None)\n",
    "            DataFrame with columns [`unique_id`, `ds`, `y`] and exogenous variables.\n",
    "            If None, a previously stored dataset is required.\n",
    "        static_df : pandas, polars or spark DataFrame, optional (default=None)\n",
    "            DataFrame with columns [`unique_id`] and static exogenous.\n",
    "        val_size : int, optional (default=0)\n",
    "            Size of validation set.\n",
    "        use_init_models : bool, optional (default=False)\n",
    "            Use initial model passed when NeuralForecast object was instantiated.\n",
    "        verbose : bool (default=False)\n",
    "            Print processing steps.\n",
    "        id_col : str (default='unique_id')\n",
    "            Column that identifies each serie.\n",
    "        time_col : str (default='ds')\n",
    "            Column that identifies each timestep, its values can be timestamps or integers.\n",
    "        target_col : str (default='y')\n",
    "            Column that contains the target.\n",
    "        distributed_config : neuralforecast.DistributedConfig\n",
    "            Configuration to use for DDP training. Currently only spark is supported.\n",
    "        prediction_intervals : PredictionIntervals, optional (default=None)\n",
    "            Configuration to calibrate prediction intervals (Conformal Prediction).            \n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        self : NeuralForecast\n",
    "            Returns `NeuralForecast` class with fitted `models`.\n",
    "        \"\"\"\n",
    "        if (df is None) and not (hasattr(self, 'dataset')):\n",
    "            raise Exception('You must pass a DataFrame or have one stored.')\n",
    "\n",
    "        # Model and datasets interactions protections\n",
    "        if (\n",
    "            any(model.early_stop_patience_steps > 0 for model in self.models)\n",
    "            and val_size == 0\n",
    "        ):\n",
    "            raise Exception('Set val_size>0 if early stopping is enabled.')\n",
    "        \n",
    "        if (val_size is not None) and (0 < val_size < self.h):\n",
    "            raise ValueError(f'val_size must be either 0 or greater than or equal to the horizon: {self.h}')\n",
    "        \n",
    "        self._cs_df: Optional[DataFrame] = None\n",
    "        self.prediction_intervals: Optional[PredictionIntervals] = None\n",
    "\n",
    "        # Process and save new dataset (in self)\n",
    "        if isinstance(df, (pd.DataFrame, pl_DataFrame)):\n",
    "            validate_freq(df[time_col], self.freq)\n",
    "            self.dataset, self.uids, self.last_dates, self.ds = self._prepare_fit(\n",
    "                df=df,\n",
    "                static_df=static_df,\n",
    "                predict_only=False,\n",
    "                id_col=id_col,\n",
    "                time_col=time_col,\n",
    "                target_col=target_col,\n",
    "            )\n",
    "            if prediction_intervals is not None:\n",
    "                self.prediction_intervals = prediction_intervals\n",
    "                self._cs_df = self._conformity_scores(\n",
    "                    df=df,\n",
    "                    id_col=id_col,\n",
    "                    time_col=time_col,\n",
    "                    target_col=target_col,\n",
    "                    static_df=static_df,\n",
    "                )\n",
    "\n",
    "        elif isinstance(df, SparkDataFrame):\n",
    "            if static_df is not None and not isinstance(static_df, SparkDataFrame):\n",
    "                raise ValueError(\n",
    "                    \"`static_df` must be a spark dataframe when `df` is a spark dataframe.\"\n",
    "                )\n",
    "            self.dataset = self._prepare_fit_distributed(\n",
    "                df=df,\n",
    "                static_df=static_df,\n",
    "                id_col=id_col,\n",
    "                time_col=time_col,\n",
    "                target_col=target_col,\n",
    "                distributed_config=distributed_config,\n",
    "            )\n",
    "\n",
    "            if prediction_intervals is not None:\n",
    "                raise NotImplementedError(\"Prediction intervals are not supported for distributed training.\")\n",
    "\n",
    "        elif isinstance(df, Sequence):\n",
    "            if not all(isinstance(val, str) for val in df):\n",
    "                raise ValueError(\"All entries in the list of files must be of type string\")        \n",
    "            self.dataset = self._prepare_fit_for_local_files(\n",
    "                files_list=df,\n",
    "                static_df=static_df,\n",
    "                id_col=id_col,\n",
    "                time_col=time_col,\n",
    "                target_col=target_col,\n",
    "            )\n",
    "            self.uids = self.dataset.indices\n",
    "            self.last_dates = self.dataset.last_times\n",
    "            \n",
    "            if prediction_intervals is not None:\n",
    "                raise NotImplementedError(\"Prediction intervals are not supported for local files.\")\n",
    "            \n",
    "        elif df is None:\n",
    "            if verbose:\n",
    "                print(\"Using stored dataset.\")\n",
    "        else:\n",
    "            raise ValueError(\n",
    "                f\"`df` must be a pandas, polars or spark DataFrame, or a list of parquet files containing the series, or `None`, got: {type(df)}\"\n",
    "            )\n",
    "\n",
    "        if val_size is not None:\n",
    "            if self.dataset.min_size < val_size:\n",
    "                warnings.warn('Validation set size is larger than the shorter time-series.')\n",
    "\n",
    "        # Recover initial model if use_init_models\n",
    "        if use_init_models:\n",
    "            self._reset_models()\n",
    "\n",
    "        for i, model in enumerate(self.models):\n",
    "            self.models[i] = model.fit(\n",
    "                self.dataset, val_size=val_size, distributed_config=distributed_config\n",
    "            )\n",
    "\n",
    "        self._fitted = True\n",
    "\n",
    "    def make_future_dataframe(self, df: Optional[DFType] = None) -> DFType:\n",
    "        \"\"\"Create a dataframe with all ids and future times in the forecasting horizon.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        df : pandas or polars DataFrame, optional (default=None)\n",
    "            DataFrame with columns [`unique_id`, `ds`, `y`] and exogenous variables.\n",
    "            Only required if this is different than the one used in the fit step.\n",
    "        \"\"\"\n",
    "        if not self._fitted:\n",
    "            raise Exception('You must fit the model first.')\n",
    "        if df is not None:\n",
    "            df = ufp.sort(df, by=[self.id_col, self.time_col])\n",
    "            last_times_by_id = ufp.group_by_agg(\n",
    "                df,\n",
    "                by=self.id_col,\n",
    "                aggs={self.time_col: 'max'},\n",
    "                maintain_order=True,\n",
    "            )\n",
    "            uids = last_times_by_id[self.id_col]\n",
    "            last_times = last_times_by_id[self.time_col]\n",
    "        else:\n",
    "            uids = self.uids\n",
    "            last_times = self.last_dates\n",
    "        return ufp.make_future_dataframe(\n",
    "            uids=uids,\n",
    "            last_times=last_times,\n",
    "            freq=self.freq,\n",
    "            h=self.h,\n",
    "            id_col=self.id_col,\n",
    "            time_col=self.time_col,\n",
    "        )\n",
    "\n",
    "    def get_missing_future(\n",
    "        self, futr_df: DFType, df: Optional[DFType] = None\n",
    "    ) -> DFType:\n",
    "        \"\"\"Get the missing ids and times combinations in `futr_df`.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        futr_df : pandas or polars DataFrame\n",
    "            DataFrame with [`unique_id`, `ds`] columns and `df`'s future exogenous.\n",
    "        df : pandas or polars DataFrame, optional (default=None)\n",
    "            DataFrame with columns [`unique_id`, `ds`, `y`] and exogenous variables.\n",
    "            Only required if this is different than the one used in the fit step.\n",
    "        \"\"\"\n",
    "        expected = self.make_future_dataframe(df)\n",
    "        ids = [self.id_col, self.time_col]\n",
    "        return ufp.anti_join(expected, futr_df[ids], on=ids)\n",
    "\n",
    "    def _get_needed_futr_exog(self):\n",
    "        futr_exogs = []\n",
    "        for m in self.models:\n",
    "            if isinstance(m, BaseAuto):\n",
    "                if isinstance(m.config, dict):  # ray\n",
    "                    exogs = m.config.get('futr_exog_list', [])\n",
    "                    if hasattr(exogs, 'categories'):  # features are being tuned, get possible values\n",
    "                        exogs = exogs.categories\n",
    "                else:   # optuna\n",
    "                    exogs = m.config(MockTrial()).get('futr_exog_list', [])\n",
    "            else:  # regular model, extract them directly\n",
    "                exogs = getattr(m, 'futr_exog_list', [])\n",
    "            \n",
    "            for exog in exogs:\n",
    "                if isinstance(exog, str):\n",
    "                    futr_exogs.append(exog)\n",
    "                else:\n",
    "                    futr_exogs.extend(exog)\n",
    "\n",
    "        return set(futr_exogs)\n",
    "    \n",
    "    def _get_needed_exog(self):\n",
    "        futr_exog = self._get_needed_futr_exog()\n",
    "\n",
    "        hist_exog = []\n",
    "        for m in self.models:\n",
    "            if isinstance(m, BaseAuto):\n",
    "                if isinstance(m.config, dict):  # ray\n",
    "                    exogs = m.config.get('hist_exog_list', [])\n",
    "                    if hasattr(exogs, 'categories'):  # features are being tuned, get possible values\n",
    "                        exogs = exogs.categories\n",
    "                else:   # optuna\n",
    "                    exogs = m.config(MockTrial()).get('hist_exog_list', [])\n",
    "            else:  # regular model, extract them directly\n",
    "                exogs = getattr(m, 'hist_exog_list', [])\n",
    "            \n",
    "            for exog in exogs:\n",
    "                if isinstance(exog, str):\n",
    "                    hist_exog.append(exog)\n",
    "                else:\n",
    "                    hist_exog.extend(exog)\n",
    "\n",
    "        return futr_exog | set(hist_exog)\n",
    "    \n",
    "    def _get_model_names(self, add_level=False) -> List[str]:\n",
    "        names: List[str] = []\n",
    "        count_names = {'model': 0}\n",
    "        for model in self.models:\n",
    "            model_name = repr(model)\n",
    "            count_names[model_name] = count_names.get(model_name, -1) + 1\n",
    "            if count_names[model_name] > 0:\n",
    "                model_name += str(count_names[model_name])\n",
    "\n",
    "            if add_level and (model.loss.outputsize_multiplier > 1 or isinstance(model.loss, (IQLoss, HuberIQLoss))):\n",
    "                continue\n",
    "\n",
    "            names.extend(model_name + n for n in model.loss.output_names)\n",
    "        return names\n",
    "\n",
    "    def _predict_distributed(\n",
    "        self,\n",
    "        df: Optional[SparkDataFrame],\n",
    "        static_df: Optional[SparkDataFrame],\n",
    "        futr_df: Optional[SparkDataFrame],\n",
    "        engine,\n",
    "    ):\n",
    "        import fugue.api as fa\n",
    "\n",
    "        def _predict(\n",
    "            df: pd.DataFrame,\n",
    "            static_cols,\n",
    "            futr_exog_cols,\n",
    "            models,\n",
    "            freq,\n",
    "            id_col,\n",
    "            time_col,\n",
    "            target_col,\n",
    "        ) -> pd.DataFrame:\n",
    "            from neuralforecast import NeuralForecast\n",
    "\n",
    "            nf = NeuralForecast(models=models, freq=freq)\n",
    "            nf.id_col = id_col\n",
    "            nf.time_col = time_col\n",
    "            nf.target_col = target_col\n",
    "            nf.scalers_ = {}\n",
    "            nf._fitted = True\n",
    "            if futr_exog_cols:\n",
    "                # if we have futr_exog we'll have extra rows with the future values\n",
    "                futr_rows = df[target_col].isnull()\n",
    "                futr_df = df.loc[futr_rows, [self.id_col, self.time_col] + futr_exog_cols].copy()\n",
    "                df = df[~futr_rows].copy()\n",
    "            else:\n",
    "                futr_df = None\n",
    "            if static_cols:\n",
    "                static_df = df[[self.id_col] + static_cols].groupby(self.id_col, observed=True).head(1)\n",
    "                df = df.drop(columns=static_cols)\n",
    "            else:\n",
    "                static_df = None\n",
    "            return nf.predict(df=df, static_df=static_df, futr_df=futr_df)\n",
    "\n",
    "        # df\n",
    "        if isinstance(df, SparkDataFrame):\n",
    "            repartition = True\n",
    "        else:\n",
    "            if engine is None:\n",
    "                raise ValueError(\"engine is required for distributed inference\")\n",
    "            df = engine.read.parquet(*self.dataset.files)\n",
    "            # we save the datataset with partitioning\n",
    "            repartition = False\n",
    "\n",
    "        # static\n",
    "        static_cols = set(chain.from_iterable(getattr(m, 'stat_exog_list', []) for m in self.models))\n",
    "        if static_df is not None:\n",
    "            if not isinstance(static_df, SparkDataFrame):\n",
    "                raise ValueError(\n",
    "                    \"`static_df` must be a spark dataframe when `df` is a spark dataframe \"\n",
    "                    \"or the models were trained in a distributed setting.\\n\"\n",
    "                    \"You can also provide local dataframes (pandas or polars) as `df` and `static_df`.\"\n",
    "                )\n",
    "            missing_static = static_cols - set(static_df.columns)\n",
    "            if missing_static:\n",
    "                raise ValueError(\n",
    "                    f\"The following static columns are missing from the static_df: {missing_static}\"\n",
    "                )\n",
    "            # join is supposed to preserve the partitioning\n",
    "            df = df.join(static_df, on=[self.id_col], how=\"left\")\n",
    "\n",
    "        # exog\n",
    "        if futr_df is not None:\n",
    "            if not isinstance(futr_df, SparkDataFrame):\n",
    "                raise ValueError(\n",
    "                    \"`futr_df` must be a spark dataframe when `df` is a spark dataframe \"\n",
    "                    \"or the models were trained in a distributed setting.\\n\"\n",
    "                    \"You can also provide local dataframes (pandas or polars) as `df` and `futr_df`.\"\n",
    "                )\n",
    "            if self.target_col in futr_df.columns:\n",
    "                raise ValueError(\"`futr_df` must not contain the target column.\")\n",
    "            # df has the statics, historic exog and target at this point, futr_df doesnt\n",
    "            df = df.unionByName(futr_df, allowMissingColumns=True)\n",
    "            # union doesn't guarantee preserving the partitioning\n",
    "            repartition = True\n",
    "\n",
    "        if repartition:\n",
    "            df = df.repartitionByRange(df.rdd.getNumPartitions(), self.id_col)    \n",
    "\n",
    "        # predict\n",
    "        base_schema = fa.get_schema(df).extract([self.id_col, self.time_col])\n",
    "        models_schema = {model: 'float' for model in self._get_model_names()}\n",
    "        return fa.transform(\n",
    "            df=df,\n",
    "            using=_predict,\n",
    "            schema=base_schema.append(models_schema),\n",
    "            params=dict(\n",
    "                static_cols=list(static_cols),\n",
    "                futr_exog_cols=list(self._get_needed_futr_exog()),\n",
    "                models=self.models,\n",
    "                freq=self.freq,\n",
    "                id_col=self.id_col,\n",
    "                time_col=self.time_col,\n",
    "                target_col=self.target_col,\n",
    "            ),\n",
    "        )\n",
    "\n",
    "    def predict(\n",
    "        self,\n",
    "        df: Optional[Union[DataFrame, SparkDataFrame]] = None,\n",
    "        static_df: Optional[Union[DataFrame, SparkDataFrame]] = None,\n",
    "        futr_df: Optional[Union[DataFrame, SparkDataFrame]] = None,\n",
    "        verbose: bool = False,\n",
    "        engine = None,\n",
    "        level: Optional[List[Union[int, float]]] = None,\n",
    "        quantiles: Optional[List[float]] = None,\n",
    "        **data_kwargs\n",
    "    ):\n",
    "        \"\"\"Predict with core.NeuralForecast.\n",
    "\n",
    "        Use stored fitted `models` to predict large set of time series from DataFrame `df`.        \n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        df : pandas, polars or spark DataFrame, optional (default=None)\n",
    "            DataFrame with columns [`unique_id`, `ds`, `y`] and exogenous variables.\n",
    "            If a DataFrame is passed, it is used to generate forecasts.\n",
    "        static_df : pandas, polars or spark DataFrame, optional (default=None)\n",
    "            DataFrame with columns [`unique_id`] and static exogenous.\n",
    "        futr_df : pandas, polars or spark DataFrame, optional (default=None)\n",
    "            DataFrame with [`unique_id`, `ds`] columns and `df`'s future exogenous.\n",
    "        verbose : bool (default=False)\n",
    "            Print processing steps.\n",
    "        engine : spark session\n",
    "            Distributed engine for inference. Only used if df is a spark dataframe or if fit was called on a spark dataframe.\n",
    "        level : list of ints or floats, optional (default=None)\n",
    "            Confidence levels between 0 and 100.\n",
    "        quantiles : list of floats, optional (default=None)\n",
    "            Alternative to level, target quantiles to predict.\n",
    "        data_kwargs : kwargs\n",
    "            Extra arguments to be passed to the dataset within each model.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        fcsts_df : pandas or polars DataFrame\n",
    "            DataFrame with insample `models` columns for point predictions and probabilistic\n",
    "            predictions for all fitted `models`.    \n",
    "        \"\"\"\n",
    "        if df is None and not hasattr(self, 'dataset'):\n",
    "            raise Exception('You must pass a DataFrame or have one stored.')\n",
    "\n",
    "        if not self._fitted:\n",
    "            raise Exception(\"You must fit the model before predicting.\")\n",
    "        \n",
    "        quantiles_ = None\n",
    "        level_ = None\n",
    "        has_level = False   \n",
    "        if level is not None:\n",
    "            has_level = True\n",
    "            if quantiles is not None:\n",
    "                raise ValueError(\"You can't set both level and quantiles.\")\n",
    "            level_ = sorted(list(set(level)))\n",
    "            quantiles_ = level_to_quantiles(level_)\n",
    "        \n",
    "        if quantiles is not None:\n",
    "            if level is not None:\n",
    "                raise ValueError(\"You can't set both level and quantiles.\")            \n",
    "            quantiles_ = sorted(list(set(quantiles)))\n",
    "            level_ = quantiles_to_level(quantiles_)\n",
    "\n",
    "        needed_futr_exog = self._get_needed_futr_exog()\n",
    "        if needed_futr_exog:\n",
    "            if futr_df is None:\n",
    "                raise ValueError(\n",
    "                    f'Models require the following future exogenous features: {needed_futr_exog}. '\n",
    "                    'Please provide them through the `futr_df` argument.'\n",
    "                )\n",
    "            else:\n",
    "                missing = needed_futr_exog - set(futr_df.columns)\n",
    "                if missing:\n",
    "                    raise ValueError(f'The following features are missing from `futr_df`: {missing}')\n",
    "\n",
    "        # distributed df or NeuralForecast instance was trained with a distributed input and no df is provided\n",
    "        # we assume the user wants to perform distributed inference as well\n",
    "        is_files_dataset = isinstance(getattr(self, 'dataset', None), _FilesDataset)\n",
    "        is_dataset_local_files = isinstance(getattr(self, 'dataset', None), LocalFilesTimeSeriesDataset)\n",
    "        if isinstance(df, SparkDataFrame) or (df is None and is_files_dataset):\n",
    "            return self._predict_distributed(\n",
    "                df=df,\n",
    "                static_df=static_df,\n",
    "                futr_df=futr_df,\n",
    "                engine=engine,\n",
    "            )\n",
    "        \n",
    "        if is_dataset_local_files and df is None:\n",
    "            raise ValueError(\n",
    "                \"When the model has been trained on a dataset that is split between multiple files, you must pass in a specific dataframe for prediciton.\"\n",
    "            )\n",
    "        \n",
    "        # Process new dataset but does not store it.\n",
    "        if df is not None:\n",
    "            validate_freq(df[self.time_col], self.freq)\n",
    "            dataset, uids, last_dates, _ = self._prepare_fit(\n",
    "                df=df,\n",
    "                static_df=static_df,\n",
    "                predict_only=True,\n",
    "                id_col=self.id_col,\n",
    "                time_col=self.time_col,\n",
    "                target_col=self.target_col,\n",
    "            )\n",
    "        else:\n",
    "            dataset = self.dataset\n",
    "            uids = self.uids\n",
    "            last_dates = self.last_dates\n",
    "            if verbose: print('Using stored dataset.')\n",
    "  \n",
    "\n",
    "        # Placeholder dataframe for predictions with unique_id and ds\n",
    "        fcsts_df = ufp.make_future_dataframe(\n",
    "            uids=uids,\n",
    "            last_times=last_dates,\n",
    "            freq=self.freq,\n",
    "            h=self.h,\n",
    "            id_col=self.id_col,\n",
    "            time_col=self.time_col,\n",
    "        )\n",
    "\n",
    "        # Update and define new forecasting dataset\n",
    "        if futr_df is None:\n",
    "            futr_df = fcsts_df\n",
    "        else:\n",
    "            futr_orig_rows = futr_df.shape[0]\n",
    "            futr_df = ufp.join(futr_df, fcsts_df, on=[self.id_col, self.time_col])\n",
    "            if futr_df.shape[0] < fcsts_df.shape[0]:\n",
    "                if df is None:\n",
    "                    expected_cmd = 'make_future_dataframe()'\n",
    "                    missing_cmd = 'get_missing_future(futr_df)'\n",
    "                else:\n",
    "                    expected_cmd = 'make_future_dataframe(df)'\n",
    "                    missing_cmd = 'get_missing_future(futr_df, df)'\n",
    "                raise ValueError(\n",
    "                    'There are missing combinations of ids and times in `futr_df`.\\n'\n",
    "                    f'You can run the `{expected_cmd}` method to get the expected combinations or '\n",
    "                    f'the `{missing_cmd}` method to get the missing combinations.'\n",
    "                )\n",
    "            if futr_orig_rows > futr_df.shape[0]:\n",
    "                dropped_rows = futr_orig_rows - futr_df.shape[0]\n",
    "                warnings.warn(\n",
    "                    f'Dropped {dropped_rows:,} unused rows from `futr_df`.'\n",
    "                )\n",
    "            if any(ufp.is_none(futr_df[col]).any() for col in needed_futr_exog):\n",
    "                raise ValueError('Found null values in `futr_df`')\n",
    "        futr_dataset = dataset.align(\n",
    "            futr_df,\n",
    "            id_col=self.id_col,\n",
    "            time_col=self.time_col,\n",
    "            target_col=self.target_col,\n",
    "        )\n",
    "        self._scalers_transform(futr_dataset)\n",
    "        dataset = dataset.append(futr_dataset)\n",
    "      \n",
    "        fcsts, cols = self._generate_forecasts(dataset=dataset, uids=uids, quantiles_=quantiles_, level_=level_, has_level=has_level, **data_kwargs)\n",
    "        \n",
    "        if self.scalers_:\n",
    "            indptr = np.append(0, np.full(len(uids), self.h).cumsum())\n",
    "            fcsts = self._scalers_target_inverse_transform(fcsts, indptr)\n",
    "\n",
    "        # Declare predictions pd.DataFrame\n",
    "        if isinstance(fcsts_df, pl_DataFrame):\n",
    "            fcsts = pl_DataFrame(dict(zip(cols, fcsts.T)))\n",
    "        else:\n",
    "            fcsts = pd.DataFrame(fcsts, columns=cols)\n",
    "        fcsts_df = ufp.horizontal_concat([fcsts_df, fcsts])\n",
    "\n",
    "        return fcsts_df\n",
    "\n",
    "    def _reset_models(self):\n",
    "        self.models = [deepcopy(model) for model in self.models_init]\n",
    "        if self._fitted:\n",
    "            print('WARNING: Deleting previously fitted models.')        \n",
    "    \n",
    "    def _no_refit_cross_validation(\n",
    "        self,\n",
    "        df: Optional[DataFrame],\n",
    "        static_df: Optional[DataFrame],\n",
    "        n_windows: int,\n",
    "        step_size: int,\n",
    "        val_size: Optional[int], \n",
    "        test_size: int,\n",
    "        verbose: bool,\n",
    "        id_col: str,\n",
    "        time_col: str,\n",
    "        target_col: str,\n",
    "        **data_kwargs\n",
    "    ) -> DataFrame:\n",
    "        if (df is None) and not (hasattr(self, 'dataset')):\n",
    "            raise Exception('You must pass a DataFrame or have one stored.')\n",
    "\n",
    "        # Process and save new dataset (in self)\n",
    "        if df is not None:\n",
    "            validate_freq(df[time_col], self.freq)\n",
    "            self.dataset, self.uids, self.last_dates, self.ds = self._prepare_fit(\n",
    "                df=df,\n",
    "                static_df=static_df,\n",
    "                predict_only=False,\n",
    "                id_col=id_col,\n",
    "                time_col=time_col,\n",
    "                target_col=target_col,\n",
    "            )\n",
    "        else:\n",
    "            if verbose: print('Using stored dataset.')\n",
    "\n",
    "        if val_size is not None:\n",
    "            if self.dataset.min_size < (val_size+test_size):\n",
    "                warnings.warn('Validation and test sets are larger than the shorter time-series.')\n",
    "\n",
    "        fcsts_df = ufp.cv_times(\n",
    "            times=self.ds,\n",
    "            uids=self.uids,\n",
    "            indptr=self.dataset.indptr,\n",
    "            h=self.h,\n",
    "            test_size=test_size,\n",
    "            step_size=step_size,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "        )\n",
    "        # the cv_times is sorted by window and then id\n",
    "        fcsts_df = ufp.sort(fcsts_df, [id_col, 'cutoff', time_col])\n",
    "\n",
    "        fcsts_list: List = []\n",
    "        for model in self.models:\n",
    "            if self._add_level and (model.loss.outputsize_multiplier > 1 or isinstance(model.loss, (IQLoss, HuberIQLoss))):\n",
    "                continue\n",
    "\n",
    "            model.fit(dataset=self.dataset,\n",
    "                        val_size=val_size, \n",
    "                        test_size=test_size)\n",
    "            model_fcsts = model.predict(self.dataset, step_size=step_size, **data_kwargs)\n",
    "\n",
    "            # Append predictions in memory placeholder\n",
    "            fcsts_list.append(model_fcsts)\n",
    "\n",
    "        fcsts = np.concatenate(fcsts_list, axis=-1)\n",
    "        # we may have allocated more space than needed\n",
    "        # each serie can produce at most (serie.size - 1) // self.h CV windows\n",
    "        effective_sizes = ufp.counts_by_id(fcsts_df, id_col)['counts'].to_numpy()\n",
    "        needs_trim = effective_sizes.sum() != fcsts.shape[0]\n",
    "        if self.scalers_ or needs_trim:\n",
    "            indptr = np.arange(\n",
    "                0,\n",
    "                n_windows * self.h * (self.dataset.n_groups + 1),\n",
    "                n_windows * self.h,\n",
    "                dtype=np.int32,\n",
    "            )\n",
    "            if self.scalers_:\n",
    "                fcsts = self._scalers_target_inverse_transform(fcsts, indptr)\n",
    "            if needs_trim:\n",
    "                # we keep only the effective samples of each serie from the cv results\n",
    "                trimmed = np.empty_like(\n",
    "                    fcsts, shape=(effective_sizes.sum(), fcsts.shape[1])\n",
    "                )\n",
    "                cv_indptr = np.append(0, effective_sizes).cumsum(dtype=np.int32)\n",
    "                for i in range(fcsts.shape[1]):\n",
    "                    ga = GroupedArray(fcsts[:, i], indptr)\n",
    "                    trimmed[:, i] = ga._tails(cv_indptr)\n",
    "                fcsts = trimmed\n",
    "\n",
    "        self._fitted = True\n",
    "\n",
    "        # Add predictions to forecasts DataFrame\n",
    "        cols = self._get_model_names(add_level=self._add_level)\n",
    "        if isinstance(self.uids, pl_Series):\n",
    "            fcsts = pl_DataFrame(dict(zip(cols, fcsts.T)))\n",
    "        else:\n",
    "            fcsts = pd.DataFrame(fcsts, columns=cols)\n",
    "        fcsts_df = ufp.horizontal_concat([fcsts_df, fcsts])\n",
    "\n",
    "        # Add original input df's y to forecasts DataFrame    \n",
    "        return ufp.join(\n",
    "            fcsts_df,\n",
    "            df[[id_col, time_col, target_col]],\n",
    "            how='left',\n",
    "            on=[id_col, time_col],\n",
    "        )  \n",
    "\n",
    "    def cross_validation(\n",
    "        self,\n",
    "        df: Optional[DataFrame] = None,\n",
    "        static_df: Optional[DataFrame] = None,\n",
    "        n_windows: int = 1,\n",
    "        step_size: int = 1,\n",
    "        val_size: Optional[int] = 0, \n",
    "        test_size: Optional[int] = None,\n",
    "        use_init_models: bool = False,\n",
    "        verbose: bool = False,\n",
    "        refit: Union[bool, int] = False,\n",
    "        id_col: str = 'unique_id',\n",
    "        time_col: str = 'ds',\n",
    "        target_col: str = 'y',\n",
    "        prediction_intervals: Optional[PredictionIntervals] = None,\n",
    "        level: Optional[List[Union[int, float]]] = None,\n",
    "        quantiles: Optional[List[float]] = None,\n",
    "        **data_kwargs\n",
    "    ) -> DataFrame:\n",
    "        \"\"\"Temporal Cross-Validation with core.NeuralForecast.\n",
    "\n",
    "        `core.NeuralForecast`'s cross-validation efficiently fits a list of NeuralForecast \n",
    "        models through multiple windows, in either chained or rolled manner.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        df : pandas or polars DataFrame, optional (default=None)\n",
    "            DataFrame with columns [`unique_id`, `ds`, `y`] and exogenous variables.\n",
    "            If None, a previously stored dataset is required.\n",
    "        static_df : pandas or polars DataFrame, optional (default=None)\n",
    "            DataFrame with columns [`unique_id`] and static exogenous.\n",
    "        n_windows : int (default=1)\n",
    "            Number of windows used for cross validation.\n",
    "        step_size : int (default=1)\n",
    "            Step size between each window.\n",
    "        val_size : int, optional (default=None)\n",
    "            Length of validation size. If passed, set `n_windows=None`.\n",
    "        test_size : int, optional (default=None)\n",
    "            Length of test size. If passed, set `n_windows=None`.\n",
    "        use_init_models : bool, option (default=False)\n",
    "            Use initial model passed when object was instantiated.\n",
    "        verbose : bool (default=False)\n",
    "            Print processing steps.\n",
    "        refit : bool or int (default=False)\n",
    "            Retrain model for each cross validation window.\n",
    "            If False, the models are trained at the beginning and then used to predict each window.\n",
    "            If positive int, the models are retrained every `refit` windows.\n",
    "        id_col : str (default='unique_id')\n",
    "            Column that identifies each serie.\n",
    "        time_col : str (default='ds')\n",
    "            Column that identifies each timestep, its values can be timestamps or integers.\n",
    "        target_col : str (default='y')\n",
    "            Column that contains the target.            \n",
    "        prediction_intervals : PredictionIntervals, optional (default=None)\n",
    "            Configuration to calibrate prediction intervals (Conformal Prediction).            \n",
    "        level : list of ints or floats, optional (default=None)\n",
    "            Confidence levels between 0 and 100.\n",
    "        quantiles : list of floats, optional (default=None)\n",
    "            Alternative to level, target quantiles to predict.\n",
    "        data_kwargs : kwargs\n",
    "            Extra arguments to be passed to the dataset within each model.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        fcsts_df : pandas or polars DataFrame\n",
    "            DataFrame with insample `models` columns for point predictions and probabilistic\n",
    "            predictions for all fitted `models`.    \n",
    "        \"\"\"\n",
    "        h = self.h\n",
    "        if n_windows is None and test_size is None:\n",
    "            raise Exception('you must define `n_windows` or `test_size`.')            \n",
    "        if test_size is None:\n",
    "            test_size = h + step_size * (n_windows - 1)\n",
    "        elif n_windows is None:\n",
    "            if (test_size - h) % step_size:\n",
    "                raise Exception('`test_size - h` should be module `step_size`')\n",
    "            n_windows = int((test_size - h) / step_size) + 1\n",
    "        else:\n",
    "            raise Exception('you must define `n_windows` or `test_size` but not both')    \n",
    "\n",
    "        # Recover initial model if use_init_models.\n",
    "        if use_init_models:\n",
    "            self._reset_models()\n",
    "\n",
    "        # Checks for prediction intervals\n",
    "        if prediction_intervals is not None:\n",
    "            if level is None and quantiles is None:\n",
    "                raise Exception('When passing prediction_intervals you need to set the level or quantiles argument.')  \n",
    "            if not refit:\n",
    "                raise Exception('Passing prediction_intervals is only supported with refit=True.')  \n",
    "\n",
    "        if level is not None and quantiles is not None:\n",
    "            raise ValueError(\"You can't set both level and quantiles argument.\")\n",
    "        \n",
    "        if not refit:\n",
    "\n",
    "            return self._no_refit_cross_validation(\n",
    "                df=df,\n",
    "                static_df=static_df,\n",
    "                n_windows=n_windows,\n",
    "                step_size=step_size,\n",
    "                val_size=val_size,\n",
    "                test_size=test_size,\n",
    "                verbose=verbose,\n",
    "                id_col=id_col,\n",
    "                time_col=time_col,\n",
    "                target_col=target_col,\n",
    "                **data_kwargs\n",
    "            )\n",
    "        if df is None:\n",
    "            raise ValueError('Must specify `df` with `refit!=False`.')\n",
    "        validate_freq(df[time_col], self.freq)\n",
    "        splits = ufp.backtest_splits(\n",
    "            df,\n",
    "            n_windows=n_windows,\n",
    "            h=self.h,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "            freq=self.freq,\n",
    "            step_size=step_size,\n",
    "            input_size=None,\n",
    "        )\n",
    "        results = []\n",
    "        for i_window, (cutoffs, train, test) in enumerate(splits):\n",
    "            should_fit = i_window == 0 or (refit > 0 and i_window % refit == 0)\n",
    "            if should_fit:\n",
    "                self.fit(\n",
    "                    df=train,\n",
    "                    static_df=static_df,\n",
    "                    val_size=val_size,\n",
    "                    use_init_models=False,\n",
    "                    verbose=verbose,\n",
    "                    id_col=id_col,\n",
    "                    time_col=time_col,\n",
    "                    target_col=target_col,\n",
    "                    prediction_intervals=prediction_intervals,                                     \n",
    "                )\n",
    "                predict_df: Optional[DataFrame] = None\n",
    "            else:\n",
    "                predict_df = train\n",
    "            needed_futr_exog = self._get_needed_futr_exog()\n",
    "            if needed_futr_exog:\n",
    "                futr_df: Optional[DataFrame] = test\n",
    "            else:\n",
    "                futr_df = None\n",
    "            preds = self.predict(\n",
    "                df=predict_df,\n",
    "                static_df=static_df,\n",
    "                futr_df=futr_df,\n",
    "                verbose=verbose,\n",
    "                level=level,\n",
    "                quantiles=quantiles,\n",
    "                **data_kwargs\n",
    "            )\n",
    "            preds = ufp.join(preds, cutoffs, on=id_col, how='left')\n",
    "            fold_result = ufp.join(\n",
    "                preds, test[[id_col, time_col, target_col]], on=[id_col, time_col]\n",
    "            )\n",
    "            results.append(fold_result)\n",
    "        out = ufp.vertical_concat(results, match_categories=False)\n",
    "        out = ufp.drop_index_if_pandas(out)\n",
    "        # match order of cv with no refit\n",
    "        first_out_cols = [id_col, time_col, \"cutoff\"]\n",
    "        remaining_cols = [\n",
    "            c for c in out.columns if c not in first_out_cols + [target_col]\n",
    "        ]\n",
    "        cols_order = first_out_cols + remaining_cols + [target_col]\n",
    "        return ufp.sort(out[cols_order], by=[id_col, 'cutoff', time_col])\n",
    "        \n",
    "    def predict_insample(self, step_size: int = 1, level: Optional[List[Union[int, float]]] = None,\n",
    "        quantiles: Optional[List[float]] = None):\n",
    "        \"\"\"Predict insample with core.NeuralForecast.\n",
    "\n",
    "        `core.NeuralForecast`'s `predict_insample` uses stored fitted `models`\n",
    "        to predict historic values of a time series from the stored dataframe.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        step_size : int (default=1)\n",
    "            Step size between each window.\n",
    "        level : list of ints or floats, optional (default=None)\n",
    "            Confidence levels between 0 and 100.\n",
    "        quantiles : list of floats, optional (default=None)\n",
    "            Alternative to level, target quantiles to predict.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        fcsts_df : pandas.DataFrame\n",
    "            DataFrame with insample predictions for all fitted `models`.    \n",
    "        \"\"\"\n",
    "        if not self._fitted:\n",
    "            raise Exception('The models must be fitted first with `fit` or `cross_validation`.')\n",
    "        test_size = self.models[0].get_test_size()\n",
    "        \n",
    "        quantiles_ = None\n",
    "        level_ = None\n",
    "        has_level = False\n",
    "        if level is not None:\n",
    "            has_level = True\n",
    "            if quantiles is not None:\n",
    "                raise ValueError(\"You can't set both level and quantiles.\")\n",
    "            level_ = sorted(list(set(level)))\n",
    "            quantiles_ = level_to_quantiles(level_)\n",
    "            if self._cs_df is not None:\n",
    "                raise NotImplementedError(\"One or more models has been trained with conformal prediction intervals. They are not supported for insample predictions. Set level=None\")\n",
    "\n",
    "        if quantiles is not None:\n",
    "            if level is not None:\n",
    "                raise ValueError(\"You can't set both level and quantiles.\")\n",
    "            quantiles_ = sorted(list(set(quantiles)))\n",
    "            level_ = quantiles_to_level(quantiles_)\n",
    "            if self._cs_df is not None:\n",
    "                raise NotImplementedError(\"One or more models has been trained with conformal prediction intervals. They are not supported for insample predictions. Set quantiles=None\")\n",
    "        \n",
    "        for model in self.models:\n",
    "            if model.MULTIVARIATE:\n",
    "                raise NotImplementedError(\n",
    "                    f\"Model {model} is multivariate. Insample predictions are not supported for multivariate models.\"\n",
    "                )\n",
    "\n",
    "        # Process each series separately\n",
    "        fcsts_dfs = []\n",
    "        trimmed_datasets = []\n",
    "        \n",
    "        for i in range(self.dataset.n_groups):\n",
    "            # Calculate series-specific length and offset\n",
    "            series_length = self.dataset.indptr[i + 1] - self.dataset.indptr[i]\n",
    "            _, forefront_offset = np.divmod((series_length - test_size - self.h), step_size)\n",
    "            \n",
    "            if test_size > 0 or forefront_offset > 0:\n",
    "                # Create single-series dataset\n",
    "                series_dataset = TimeSeriesDataset(\n",
    "                    temporal=self.dataset.temporal[self.dataset.indptr[i]:self.dataset.indptr[i + 1]],\n",
    "                    temporal_cols=self.dataset.temporal_cols,\n",
    "                    indptr=np.array([0, series_length]),\n",
    "                    y_idx=self.dataset.y_idx\n",
    "                )\n",
    "                \n",
    "                # Trim the series\n",
    "                trimmed_series = TimeSeriesDataset.trim_dataset(\n",
    "                    dataset=series_dataset,\n",
    "                    right_trim=test_size,\n",
    "                    left_trim=forefront_offset\n",
    "                )\n",
    "                \n",
    "                new_idxs = np.arange(\n",
    "                    self.dataset.indptr[i] + forefront_offset,\n",
    "                    self.dataset.indptr[i + 1] - test_size\n",
    "                )\n",
    "                times = self.ds[new_idxs]\n",
    "            else:\n",
    "                trimmed_series = TimeSeriesDataset(\n",
    "                    temporal=self.dataset.temporal[self.dataset.indptr[i]:self.dataset.indptr[i + 1]],\n",
    "                    temporal_cols=self.dataset.temporal_cols,\n",
    "                    indptr=np.array([0, series_length]),\n",
    "                    y_idx=self.dataset.y_idx\n",
    "                )\n",
    "                times = self.ds[self.dataset.indptr[i]:self.dataset.indptr[i + 1]]\n",
    "            \n",
    "            series_fcsts_df = _insample_times(\n",
    "                times=times,\n",
    "                uids=self.uids[i:i+1],\n",
    "                indptr=trimmed_series.indptr,\n",
    "                h=self.h,\n",
    "                freq=self.freq,\n",
    "                step_size=step_size,\n",
    "                id_col=self.id_col,\n",
    "                time_col=self.time_col,\n",
    "            )\n",
    "            \n",
    "            fcsts_dfs.append(series_fcsts_df)\n",
    "            trimmed_datasets.append(trimmed_series)\n",
    "\n",
    "        # Combine all series forecasts DataFrames\n",
    "        fcsts_df = ufp.vertical_concat(fcsts_dfs)\n",
    "        \n",
    "        h_backup = self.h\n",
    "        fcst_list = []\n",
    "        # Generate predictions for each dataset\n",
    "        for i, trimmed_dataset in enumerate(trimmed_datasets):\n",
    "            # Set test size to current series length\n",
    "            self.h = trimmed_dataset.max_size\n",
    "            fcsts, cols = self._generate_forecasts(\n",
    "                dataset=trimmed_dataset,\n",
    "                uids=self.uids[i:i+1],\n",
    "                quantiles_=quantiles_,\n",
    "                level_=level_,\n",
    "                has_level=has_level,\n",
    "                step_size=step_size,\n",
    "            )\n",
    "            fcst_list.append(fcsts)\n",
    "        \n",
    "        fcsts = np.vstack(fcst_list)       \n",
    "        self.h = h_backup\n",
    "\n",
    "        # Add original y values\n",
    "        original_y = {\n",
    "            self.id_col: ufp.repeat(self.uids, np.diff(self.dataset.indptr)),\n",
    "            self.time_col: self.ds,\n",
    "            self.target_col: self.dataset.temporal[:, 0].numpy(),\n",
    "        }\n",
    "\n",
    "        # Declare predictions pd.DataFrame\n",
    "        if isinstance(fcsts_df, pl_DataFrame):\n",
    "            fcsts = pl_DataFrame(dict(zip(cols, fcsts.T)))\n",
    "            Y_df = pl_DataFrame(original_y)\n",
    "        else:\n",
    "            fcsts = pd.DataFrame(fcsts, columns=cols)\n",
    "            Y_df = pd.DataFrame(original_y).reset_index(drop=True)\n",
    "\n",
    "        fcsts_df = ufp.horizontal_concat([fcsts_df, fcsts])\n",
    "        fcsts_df = ufp.join(fcsts_df, Y_df, how='left', on=[self.id_col, self.time_col])\n",
    "\n",
    "        if self.scalers_:\n",
    "            sizes = ufp.counts_by_id(fcsts_df, self.id_col)['counts'].to_numpy()\n",
    "            indptr = np.append(0, sizes.cumsum())\n",
    "            invert_cols = cols + [self.target_col]\n",
    "            fcsts_df[invert_cols] = self._scalers_target_inverse_transform(\n",
    "                fcsts_df[invert_cols].to_numpy(),\n",
    "                indptr\n",
    "            )\n",
    "\n",
    "        return fcsts_df\n",
    "\n",
    "    # Save list of models with pytorch lightning save_checkpoint function\n",
    "    def save(self, path: str, model_index: Optional[List]=None, save_dataset: bool=True, overwrite: bool=False):\n",
    "        \"\"\"Save NeuralForecast core class.\n",
    "\n",
    "        `core.NeuralForecast`'s method to save current status of models, dataset, and configuration.\n",
    "        Note that by default the `models` are not saving training checkpoints to save disk memory,\n",
    "        to get them change the individual model `**trainer_kwargs` to include `enable_checkpointing=True`.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        path : str\n",
    "            Directory to save current status.\n",
    "        model_index : list, optional (default=None)\n",
    "            List to specify which models from list of self.models to save.\n",
    "        save_dataset : bool (default=True)\n",
    "            Whether to save dataset or not.\n",
    "        overwrite : bool (default=False)\n",
    "            Whether to overwrite files or not.\n",
    "        \"\"\"\n",
    "        # Standarize path without '/'\n",
    "        if path[-1] == '/':\n",
    "            path = path[:-1]\n",
    "\n",
    "        # Model index list\n",
    "        if model_index is None:\n",
    "            model_index = list(range(len(self.models)))\n",
    "\n",
    "        fs, _, _ = fsspec.get_fs_token_paths(path)\n",
    "        if not fs.exists(path):\n",
    "            fs.makedirs(path)\n",
    "        else:\n",
    "            # Check if directory is empty to protect overwriting files\n",
    "            files = fs.ls(path)\n",
    "\n",
    "            # Checking if the list is empty or not\n",
    "            if files:\n",
    "                if not overwrite:\n",
    "                    raise Exception('Directory is not empty. Set `overwrite=True` to overwrite files.')\n",
    "                else:\n",
    "                    fs.rm(path, recursive=True)\n",
    "                    fs.mkdir(path)\n",
    "\n",
    "        # Save models\n",
    "        count_names = {'model': 0}\n",
    "        alias_to_model = {}\n",
    "        for i, model in enumerate(self.models):\n",
    "            # Skip model if not in list\n",
    "            if i not in model_index:\n",
    "                continue\n",
    "\n",
    "            model_name = repr(model)\n",
    "            if model.__class__.__name__.lower() in MODEL_FILENAME_DICT:\n",
    "                model_class_name = model.__class__.__name__.lower()\n",
    "            elif model.__class__.__base__.__name__.lower() in MODEL_FILENAME_DICT:\n",
    "                model_class_name = model.__class__.__base__.__name__.lower()\n",
    "            else:\n",
    "                raise ValueError(\n",
    "                    f\"Model {model.__class__.__name__} is not supported for saving.\"\n",
    "                )\n",
    "            alias_to_model[model_name] = model_class_name\n",
    "            count_names[model_name] = count_names.get(model_name, -1) + 1\n",
    "            model.save(f\"{path}/{model_name}_{count_names[model_name]}.ckpt\")\n",
    "        with fsspec.open(f\"{path}/alias_to_model.pkl\", \"wb\") as f:\n",
    "            pickle.dump(alias_to_model, f)\n",
    "\n",
    "        # Save dataset\n",
    "        if save_dataset and hasattr(self, 'dataset'):\n",
    "            if isinstance(self.dataset, _FilesDataset):\n",
    "                raise ValueError(\n",
    "                    \"Cannot save distributed dataset.\\n\"\n",
    "                    \"You can set `save_dataset=False` and use the `df` argument in the predict method after loading \"\n",
    "                    \"this model to use it for inference.\"\n",
    "                )\n",
    "            with fsspec.open(f\"{path}/dataset.pkl\", \"wb\") as f:\n",
    "                pickle.dump(self.dataset, f)\n",
    "        elif save_dataset:\n",
    "            raise Exception('You need to have a stored dataset to save it, \\\n",
    "                             set `save_dataset=False` to skip saving dataset.')\n",
    "\n",
    "        # Save configuration and parameters\n",
    "        config_dict = {\n",
    "            \"h\": self.h,\n",
    "            \"freq\": self.freq,\n",
    "            \"_fitted\": self._fitted,\n",
    "            \"local_scaler_type\": self.local_scaler_type,\n",
    "            \"scalers_\": self.scalers_,\n",
    "            \"id_col\": self.id_col,\n",
    "            \"time_col\": self.time_col,\n",
    "            \"target_col\": self.target_col,\n",
    "        }\n",
    "        for attr in ['prediction_intervals', '_cs_df']:\n",
    "            # conformal prediction related attributes was not available < 1.7.6\n",
    "            config_dict[attr] = getattr(self, attr, None)\n",
    "            \n",
    "\n",
    "        if save_dataset:\n",
    "            config_dict.update(\n",
    "                {\n",
    "                    \"uids\": self.uids,\n",
    "                    \"last_dates\": self.last_dates,\n",
    "                    \"ds\": self.ds,\n",
    "                }\n",
    "            )\n",
    "\n",
    "        with fsspec.open(f\"{path}/configuration.pkl\", \"wb\") as f:\n",
    "            pickle.dump(config_dict, f)\n",
    "    \n",
    "    @staticmethod\n",
    "    def load(path, verbose=False, **kwargs):\n",
    "        \"\"\"Load NeuralForecast\n",
    "\n",
    "        `core.NeuralForecast`'s method to load checkpoint from path.\n",
    "\n",
    "        Parameters\n",
    "        -----------\n",
    "        path : str\n",
    "            Directory with stored artifacts.\n",
    "        kwargs\n",
    "            Additional keyword arguments to be passed to the function\n",
    "            `load_from_checkpoint`.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        result : NeuralForecast\n",
    "            Instantiated `NeuralForecast` class.\n",
    "        \"\"\"\n",
    "        # Standarize path without '/'\n",
    "        if path[-1] == \"/\":\n",
    "            path = path[:-1]\n",
    "\n",
    "        fs, _, _ = fsspec.get_fs_token_paths(path)\n",
    "        files = [f.split(\"/\")[-1] for f in fs.ls(path) if fs.isfile(f)]\n",
    "\n",
    "        # Load models\n",
    "        models_ckpt = [f for f in files if f.endswith(\".ckpt\")]\n",
    "        if len(models_ckpt) == 0:\n",
    "            raise Exception(\"No model found in directory.\")\n",
    "\n",
    "        if verbose:\n",
    "            print(10 * \"-\" + \" Loading models \" + 10 * \"-\")\n",
    "        models = []\n",
    "        try:\n",
    "            with fsspec.open(f\"{path}/alias_to_model.pkl\", \"rb\") as f:\n",
    "                alias_to_model = pickle.load(f)\n",
    "        except FileNotFoundError:\n",
    "            alias_to_model = {}\n",
    "        \n",
    "        for model in models_ckpt:\n",
    "            model_name = \"_\".join(model.split(\"_\")[:-1])\n",
    "            model_class_name = alias_to_model.get(model_name, model_name)\n",
    "            loaded_model = MODEL_FILENAME_DICT[model_class_name].load(\n",
    "                f\"{path}/{model}\", **kwargs\n",
    "            )\n",
    "            loaded_model.alias = model_name\n",
    "            models.append(loaded_model)\n",
    "            if verbose:\n",
    "                print(f\"Model {model_name} loaded.\")\n",
    "\n",
    "        if verbose:\n",
    "            print(10 * \"-\" + \" Loading dataset \" + 10 * \"-\")\n",
    "        # Load dataset\n",
    "        try:\n",
    "            with fsspec.open(f\"{path}/dataset.pkl\", \"rb\") as f:\n",
    "                dataset = pickle.load(f)\n",
    "            if verbose:\n",
    "                print(\"Dataset loaded.\")\n",
    "        except FileNotFoundError:\n",
    "            dataset = None\n",
    "            if verbose:\n",
    "                print(\"No dataset found in directory.\")\n",
    "\n",
    "        if verbose:\n",
    "            print(10 * \"-\" + \" Loading configuration \" + 10 * \"-\")\n",
    "        # Load configuration\n",
    "        try:\n",
    "            with fsspec.open(f\"{path}/configuration.pkl\", \"rb\") as f:\n",
    "                config_dict = pickle.load(f)\n",
    "            if verbose:\n",
    "                print(\"Configuration loaded.\")\n",
    "        except FileNotFoundError:\n",
    "            raise Exception(\"No configuration found in directory.\")\n",
    "\n",
    "        # in 1.6.4, `local_scaler_type` / `scalers_` lived on the dataset.\n",
    "        # in order to preserve backwards-compatibility, we check to see if these are found on the dataset\n",
    "        # in case they cannot be found in `config_dict`\n",
    "        default_scalar_type = getattr(dataset, \"local_scaler_type\", None)\n",
    "        default_scalars_ = getattr(dataset, \"scalers_\", None)\n",
    "\n",
    "        # Create NeuralForecast object\n",
    "        neuralforecast = NeuralForecast(\n",
    "            models=models,\n",
    "            freq=config_dict[\"freq\"],\n",
    "            local_scaler_type=config_dict.get(\"local_scaler_type\", default_scalar_type),\n",
    "        )\n",
    "\n",
    "        attr_to_default = {\"id_col\": \"unique_id\", \"time_col\": \"ds\", \"target_col\": \"y\"}\n",
    "        for attr, default in attr_to_default.items():\n",
    "            setattr(neuralforecast, attr, config_dict.get(attr, default))\n",
    "        # only restore attribute if available\n",
    "        for attr in [\"prediction_intervals\", \"_cs_df\"]:\n",
    "            setattr(neuralforecast, attr, config_dict.get(attr, None))\n",
    "\n",
    "        # Dataset\n",
    "        if dataset is not None:\n",
    "            neuralforecast.dataset = dataset\n",
    "            restore_attrs = [\n",
    "                \"uids\",\n",
    "                \"last_dates\",\n",
    "                \"ds\",\n",
    "            ]\n",
    "            for attr in restore_attrs:\n",
    "                setattr(neuralforecast, attr, config_dict[attr])\n",
    "\n",
    "        # Fitted flag\n",
    "        neuralforecast._fitted = config_dict[\"_fitted\"]\n",
    "\n",
    "        neuralforecast.scalers_ = config_dict.get(\"scalers_\", default_scalars_)\n",
    "\n",
    "        return neuralforecast\n",
    "    \n",
    "    def _conformity_scores(\n",
    "        self,\n",
    "        df: DataFrame,\n",
    "        id_col: str, \n",
    "        time_col: str,\n",
    "        target_col: str,\n",
    "        static_df: Optional[DataFrame],\n",
    "    ) -> DataFrame:\n",
    "        \"\"\"Compute conformity scores.\n",
    "        \n",
    "        We need at least two cross validation errors to compute\n",
    "        quantiles for prediction intervals (`n_windows=2`, specified by self.prediction_intervals).\n",
    "        \n",
    "        The exception is raised by the PredictionIntervals data class.\n",
    "\n",
    "        df: DataFrame,\n",
    "        id_col: str,\n",
    "        time_col: str,\n",
    "        target_col: str,\n",
    "        static_df: Optional[DataFrame],\n",
    "        \"\"\"\n",
    "        if self.prediction_intervals is None:\n",
    "            raise AttributeError('Please rerun the `fit` method passing a valid prediction_interval setting to compute conformity scores')\n",
    "                    \n",
    "        min_size = ufp.counts_by_id(df, id_col)['counts'].min()\n",
    "        min_samples = self.h * self.prediction_intervals.n_windows + 1\n",
    "        if min_size < min_samples:\n",
    "            raise ValueError(\n",
    "                \"Minimum required samples in each serie for the prediction intervals \"\n",
    "                f\"settings are: {min_samples}, shortest serie has: {min_size}. \"\n",
    "                \"Please reduce the number of windows, horizon or remove those series.\"\n",
    "            )\n",
    "        \n",
    "        self._add_level = True\n",
    "        cv_results = self.cross_validation(\n",
    "            df=df,\n",
    "            static_df=static_df,\n",
    "            n_windows=self.prediction_intervals.n_windows,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "            target_col=target_col,\n",
    "        )\n",
    "        self._add_level = False\n",
    "\n",
    "        kept = [time_col, id_col, 'cutoff']\n",
    "        # conformity score for each model\n",
    "        for model in self._get_model_names(add_level=True):\n",
    "            kept.append(model)\n",
    "\n",
    "            # compute absolute error for each model\n",
    "            abs_err = abs(cv_results[model] - cv_results[target_col])\n",
    "            cv_results = ufp.assign_columns(cv_results, model, abs_err)\n",
    "        dropped = list(set(cv_results.columns) - set(kept))\n",
    "        return ufp.drop_columns(cv_results, dropped)           \n",
    "    \n",
    "    def _generate_forecasts(self, dataset: TimeSeriesDataset, uids: Series, quantiles_: Optional[List[float]] = None, level_: Optional[List[Union[int, float]]] = None, has_level: Optional[bool] = False, **data_kwargs) -> np.array:\n",
    "        fcsts_list: List = []\n",
    "        cols = []\n",
    "        count_names = {'model': 0}\n",
    "        for model in self.models:\n",
    "            old_test_size = model.get_test_size()\n",
    "            model.set_test_size(self.h) # To predict h steps ahead\n",
    "            \n",
    "            # Increment model name if the same model is used more than once\n",
    "            model_name = repr(model)\n",
    "            count_names[model_name] = count_names.get(model_name, -1) + 1\n",
    "            if count_names[model_name] > 0:\n",
    "                model_name += str(count_names[model_name])\n",
    "\n",
    "            # Predict for every quantile or level if requested and the loss function supports it\n",
    "            # case 1: DistributionLoss and MixtureLosses\n",
    "            if quantiles_ is not None and not isinstance(model.loss, (IQLoss, HuberIQLoss)) and hasattr(model.loss, 'update_quantile') and callable(model.loss.update_quantile):\n",
    "                model_fcsts = model.predict(dataset=dataset, quantiles = quantiles_, **data_kwargs)\n",
    "                fcsts_list.append(model_fcsts)      \n",
    "                col_names = []\n",
    "                for i, quantile in enumerate(quantiles_):\n",
    "                    col_name = self._get_column_name(model_name, quantile, has_level)\n",
    "                    if i == 0:\n",
    "                        col_names.extend([f\"{model_name}\", col_name])\n",
    "                    else:\n",
    "                        col_names.extend([col_name])\n",
    "                if hasattr(model.loss, 'return_params') and model.loss.return_params:\n",
    "                    cols.extend(col_names + [model_name + param_name for param_name in model.loss.param_names])\n",
    "                else:\n",
    "                    cols.extend(col_names)\n",
    "            # case 2: IQLoss\n",
    "            elif quantiles_ is not None and isinstance(model.loss, (IQLoss, HuberIQLoss)):\n",
    "                # IQLoss does not give monotonically increasing quantiles, so we apply a hack: compute all quantiles, and take the quantile over the quantiles\n",
    "                quantiles_iqloss = [0.01, 0.05, 0.1, 0.2, 0.3, 0.5, 0.7, 0.8, 0.9, 0.95, 0.99]\n",
    "                fcsts_list_iqloss = []\n",
    "                for i, quantile in enumerate(quantiles_iqloss):\n",
    "                    model_fcsts = model.predict(dataset=dataset, quantiles = [quantile], **data_kwargs)               \n",
    "                    fcsts_list_iqloss.append(model_fcsts)      \n",
    "                fcsts_iqloss = np.concatenate(fcsts_list_iqloss, axis=-1)\n",
    "\n",
    "                # Get the actual requested quantiles\n",
    "                model_fcsts = np.quantile(fcsts_iqloss, quantiles_, axis=-1).T\n",
    "                fcsts_list.append(model_fcsts)      \n",
    "\n",
    "                # Get the right column names\n",
    "                col_names = []\n",
    "                for i, quantile in enumerate(quantiles_):\n",
    "                    col_name = self._get_column_name(model_name, quantile, has_level)\n",
    "                    col_names.extend([col_name])                \n",
    "                cols.extend(col_names)\n",
    "            # case 3: PointLoss via prediction intervals\n",
    "            elif quantiles_ is not None and model.loss.outputsize_multiplier == 1:\n",
    "                if self.prediction_intervals is None:\n",
    "                    raise AttributeError(\n",
    "                    f\"You have trained {model_name} with loss={type(model.loss).__name__}(). \\n\"\n",
    "                    \" You then must set `prediction_intervals` during fit to use level or quantiles during predict.\")  \n",
    "                model_fcsts = model.predict(dataset=dataset, quantiles = quantiles_, **data_kwargs)\n",
    "                prediction_interval_method = get_prediction_interval_method(self.prediction_intervals.method)\n",
    "                fcsts_with_intervals, out_cols = prediction_interval_method(\n",
    "                    model_fcsts,\n",
    "                    self._cs_df,\n",
    "                    model=model_name,\n",
    "                    level=level_ if has_level else None,\n",
    "                    cs_n_windows=self.prediction_intervals.n_windows,\n",
    "                    n_series=len(uids),\n",
    "                    horizon=self.h,\n",
    "                    quantiles=quantiles_ if not has_level else None,\n",
    "                )  \n",
    "                fcsts_list.append(fcsts_with_intervals)      \n",
    "                cols.extend([model_name] + out_cols)\n",
    "            # base case: quantiles or levels are not supported or provided as arguments\n",
    "            else:\n",
    "                model_fcsts = model.predict(dataset=dataset, **data_kwargs)\n",
    "                fcsts_list.append(model_fcsts)\n",
    "                cols.extend(model_name + n for n in model.loss.output_names)\n",
    "            model.set_test_size(old_test_size) # Set back to original value\n",
    "        fcsts = np.concatenate(fcsts_list, axis=-1)\n",
    "\n",
    "        return fcsts, cols\n",
    "    \n",
    "    @staticmethod\n",
    "    def _get_column_name(model_name, quantile, has_level) -> str:\n",
    "        if not has_level:\n",
    "            col_name = f\"{model_name}_ql{quantile}\" \n",
    "        elif quantile < 0.5:\n",
    "            level_lo = int(round(100 - 200 * quantile))\n",
    "            col_name = f\"{model_name}-lo-{level_lo}\"\n",
    "        elif quantile > 0.5:\n",
    "            level_hi = int(round(100 - 200 * (1 - quantile)))\n",
    "            col_name = f\"{model_name}-hi-{level_hi}\"\n",
    "        else:\n",
    "            col_name = f\"{model_name}-median\"\n",
    "\n",
    "        return col_name\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d6ef366-daec-4ec6-a2ae-199c6ea39a51",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "import logging\n",
    "import warnings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ac1aa65-40a4-4909-bdfb-1439c30439b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "logging.getLogger(\"pytorch_lightning\").setLevel(logging.ERROR)\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bede563-78c0-40ee-ba76-f06f329cd772",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(NeuralForecast.fit, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f90209f6-16da-40a6-8302-1c5c2f66c619",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(NeuralForecast.predict, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19a8923a-f4f3-4e60-b9b9-a7088fc9bff5",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(NeuralForecast.cross_validation, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "355df52b",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(NeuralForecast.predict_insample, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93155738-b40f-43d3-ba76-d345bf2583d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(NeuralForecast.save, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e915796-173c-4400-812f-c6351d5df3be",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(NeuralForecast.load, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b534d29d-eecc-43ba-8468-c23305fa24a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "import tempfile\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import pytorch_lightning as pl\n",
    "\n",
    "import neuralforecast\n",
    "import optuna\n",
    "from ray import tune\n",
    "\n",
    "from neuralforecast.auto import (\n",
    "    AutoMLP, AutoNBEATS, AutoNBEATSx,\n",
    "    AutoRNN, AutoTCN, AutoDilatedRNN,\n",
    ")\n",
    "\n",
    "from neuralforecast.models.rnn import RNN\n",
    "from neuralforecast.models.tcn import TCN\n",
    "from neuralforecast.models.deepar import DeepAR\n",
    "from neuralforecast.models.dilated_rnn import DilatedRNN\n",
    "\n",
    "from neuralforecast.models.mlp import MLP\n",
    "from neuralforecast.models.nhits import NHITS\n",
    "from neuralforecast.models.nbeats import NBEATS\n",
    "from neuralforecast.models.nbeatsx import NBEATSx\n",
    "\n",
    "from neuralforecast.models.tft import TFT\n",
    "from neuralforecast.models.vanillatransformer import VanillaTransformer\n",
    "from neuralforecast.models.informer import Informer\n",
    "from neuralforecast.models.autoformer import Autoformer\n",
    "\n",
    "from neuralforecast.models.stemgnn import StemGNN\n",
    "from neuralforecast.models.tsmixer import TSMixer\n",
    "from neuralforecast.models.tsmixerx import TSMixerx\n",
    "\n",
    "from neuralforecast.losses.pytorch import (DistributionLoss,\n",
    "    GMM, HuberIQLoss, HuberMQLoss, HuberQLoss, IQLoss, MAE,\n",
    "    MQLoss, MSE, NBMM, PMM, QuantileLoss)\n",
    "\n",
    "from neuralforecast.utils import AirPassengersDF, AirPassengersPanel, AirPassengersStatic\n",
    "\n",
    "from datetime import date"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fd1507c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "AirPassengersPanel_train = AirPassengersPanel[AirPassengersPanel['ds'] < AirPassengersPanel['ds'].values[-12]].reset_index(drop=True)\n",
    "AirPassengersPanel_test = AirPassengersPanel[AirPassengersPanel['ds'] >= AirPassengersPanel['ds'].values[-12]].reset_index(drop=True)\n",
    "AirPassengersPanel_test['y'] = np.nan\n",
    "AirPassengersPanel_test['y_[lag12]'] = np.nan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2e35f8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Unitest for early stopping without val_size protection\n",
    "models = [\n",
    "    NHITS(h=12, input_size=12, max_steps=1, early_stop_patience_steps=5)\n",
    "]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "test_fail(nf.fit,\n",
    "          contains='Set val_size>0 if early stopping is enabled.',\n",
    "          args=(AirPassengersPanel_train,))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14ae4692",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test fit+cross_validation behaviour\n",
    "models = [NHITS(h=12, input_size=24, max_steps=10)]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(AirPassengersPanel_train)\n",
    "init_fcst = nf.predict()\n",
    "init_cv = nf.cross_validation(AirPassengersPanel_train, use_init_models=True)\n",
    "after_cv = nf.cross_validation(AirPassengersPanel_train, use_init_models=True)\n",
    "nf.fit(AirPassengersPanel_train, use_init_models=True)\n",
    "after_fcst = nf.predict()\n",
    "test_eq(init_cv, after_cv)\n",
    "test_eq(init_fcst, after_fcst)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fddd5459-cc9b-4fd3-b50c-493b969b83f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test cross_validation with refit\n",
    "models = [\n",
    "    NHITS(\n",
    "        h=12,\n",
    "        input_size=24,\n",
    "        max_steps=2,\n",
    "        futr_exog_list=['trend'],\n",
    "        stat_exog_list=['airline1', 'airline2']\n",
    "    )\n",
    "]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "cv_kwargs = dict(\n",
    "    df=AirPassengersPanel_train,\n",
    "    static_df=AirPassengersStatic,\n",
    "    n_windows=4,\n",
    "    use_init_models=True,\n",
    ")\n",
    "cv_res_norefit = nf.cross_validation(refit=False, **cv_kwargs)\n",
    "cutoffs = cv_res_norefit['cutoff'].unique()\n",
    "for refit in [True, 2]:\n",
    "    cv_res = nf.cross_validation(refit=refit, **cv_kwargs)\n",
    "    refit = int(refit)\n",
    "    fltr = lambda df: df['cutoff'].isin(cutoffs[:refit])\n",
    "    expected = cv_res_norefit[fltr]\n",
    "    actual = cv_res[fltr]\n",
    "    # predictions for the no-refit windows should be the same\n",
    "    pd.testing.assert_frame_equal(\n",
    "        actual.reset_index(drop=True),\n",
    "        expected.reset_index(drop=True)\n",
    "    )\n",
    "    # predictions after refit should be different\n",
    "    test_fail(\n",
    "        lambda: pd.testing.assert_frame_equal(\n",
    "            cv_res_norefit.drop(expected.index).reset_index(drop=True),\n",
    "            cv_res.drop(actual.index).reset_index(drop=True),\n",
    "        ),\n",
    "        contains='(column name=\"NHITS\") are different',\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1161197c-c0c2-4d71-9b39-701777db26e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test scaling\n",
    "models = [NHITS(h=12, input_size=24, max_steps=10)]\n",
    "models_exog = [NHITS(h=12, input_size=12, max_steps=10, hist_exog_list=['trend'], futr_exog_list=['trend'])]\n",
    "\n",
    "# fit+predict\n",
    "nf = NeuralForecast(models=models, freq='M', local_scaler_type='standard')\n",
    "nf.fit(AirPassengersPanel_train)\n",
    "scaled_fcst = nf.predict()\n",
    "# check that the forecasts are similar to the one without scaling\n",
    "np.testing.assert_allclose(\n",
    "    init_fcst['NHITS'].values,\n",
    "    scaled_fcst['NHITS'].values,\n",
    "    rtol=0.3,\n",
    ")\n",
    "# with exog\n",
    "nf = NeuralForecast(models=models_exog, freq='M', local_scaler_type='standard')\n",
    "nf.fit(AirPassengersPanel_train)\n",
    "scaled_exog_fcst = nf.predict(futr_df=AirPassengersPanel_test)\n",
    "# check that the forecasts are similar to the one without exog\n",
    "np.testing.assert_allclose(\n",
    "    scaled_fcst['NHITS'].values,\n",
    "    scaled_exog_fcst['NHITS'].values,\n",
    "    rtol=0.3,\n",
    ")\n",
    "\n",
    "# CV\n",
    "nf = NeuralForecast(models=models, freq='M', local_scaler_type='robust')\n",
    "cv_res = nf.cross_validation(AirPassengersPanel)\n",
    "# check that the forecasts are similar to the original values (originals are restored directly from the df)\n",
    "np.testing.assert_allclose(\n",
    "    cv_res['NHITS'].values,\n",
    "    cv_res['y'].values,\n",
    "    rtol=0.3,\n",
    ")\n",
    "# with exog\n",
    "nf = NeuralForecast(models=models_exog, freq='M', local_scaler_type='robust-iqr')\n",
    "cv_res_exog = nf.cross_validation(AirPassengersPanel)\n",
    "# check that the forecasts are similar to the original values (originals are restored directly from the df)\n",
    "np.testing.assert_allclose(\n",
    "    cv_res_exog['NHITS'].values,\n",
    "    cv_res_exog['y'].values,\n",
    "    rtol=0.2,\n",
    ")\n",
    "\n",
    "# fit+predict_insample\n",
    "nf = NeuralForecast(models=models, freq='M', local_scaler_type='minmax')\n",
    "nf.fit(AirPassengersPanel_train)\n",
    "insample_res = (\n",
    "    nf.predict_insample()\n",
    "    .groupby('unique_id').tail(-12) # first values aren't reliable\n",
    "    .merge(\n",
    "        AirPassengersPanel_train[['unique_id', 'ds', 'y']],\n",
    "        on=['unique_id', 'ds'],\n",
    "        how='left',\n",
    "        suffixes=('_actual', '_expected'),\n",
    "    )\n",
    ")\n",
    "# y is inverted correctly\n",
    "np.testing.assert_allclose(\n",
    "    insample_res['y_actual'].values,\n",
    "    insample_res['y_expected'].values,\n",
    "    rtol=1e-5,\n",
    ")\n",
    "# predictions are in the same scale\n",
    "np.testing.assert_allclose(\n",
    "    insample_res['NHITS'].values,\n",
    "    insample_res['y_expected'].values,\n",
    "    rtol=0.7,\n",
    ")\n",
    "# with exog\n",
    "nf = NeuralForecast(models=models_exog, freq='M', local_scaler_type='minmax')\n",
    "nf.fit(AirPassengersPanel_train)\n",
    "insample_res_exog = (\n",
    "    nf.predict_insample()\n",
    "    .groupby('unique_id').tail(-12) # first values aren't reliable\n",
    "    .merge(\n",
    "        AirPassengersPanel_train[['unique_id', 'ds', 'y']],\n",
    "        on=['unique_id', 'ds'],\n",
    "        how='left',\n",
    "        suffixes=('_actual', '_expected'),\n",
    "    )\n",
    ")\n",
    "# y is inverted correctly\n",
    "np.testing.assert_allclose(\n",
    "    insample_res_exog['y_actual'].values,\n",
    "    insample_res_exog['y_expected'].values,\n",
    "    rtol=1e-5,\n",
    ")\n",
    "# predictions are similar than without exog\n",
    "np.testing.assert_allclose(\n",
    "    insample_res['NHITS'].values,\n",
    "    insample_res_exog['NHITS'].values,\n",
    "    rtol=0.2,\n",
    ")\n",
    "\n",
    "# test boxcox\n",
    "nf = NeuralForecast(models=models, freq='M', local_scaler_type='boxcox')\n",
    "nf.fit(AirPassengersPanel_train)\n",
    "insample_res = (\n",
    "    nf.predict_insample()\n",
    "    .groupby('unique_id').tail(-12) # first values aren't reliable\n",
    "    .merge(\n",
    "        AirPassengersPanel_train[['unique_id', 'ds', 'y']],\n",
    "        on=['unique_id', 'ds'],\n",
    "        how='left',\n",
    "        suffixes=('_actual', '_expected'),\n",
    "    )\n",
    ")\n",
    "# y is inverted correctly\n",
    "np.testing.assert_allclose(\n",
    "    insample_res['y_actual'].values,\n",
    "    insample_res['y_expected'].values,\n",
    "    rtol=1e-5,\n",
    ")\n",
    "# predictions are in the same scale\n",
    "np.testing.assert_allclose(\n",
    "    insample_res['NHITS'].values,\n",
    "    insample_res['y_expected'].values,\n",
    "    rtol=0.7,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1b34d37-29b0-49d4-9c7b-cb8add7ce9cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test futr_df contents\n",
    "models = [NHITS(h=6, input_size=24, max_steps=10, hist_exog_list=['trend'], futr_exog_list=['trend'])]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(AirPassengersPanel_train)\n",
    "# not enough rows in futr_df raises an error\n",
    "test_fail(lambda: nf.predict(futr_df=AirPassengersPanel_test.head()), contains='There are missing combinations')\n",
    "# extra rows issues a warning\n",
    "with warnings.catch_warnings(record=True) as issued_warnings:\n",
    "    warnings.simplefilter('always', UserWarning)\n",
    "    nf.predict(futr_df=AirPassengersPanel_test)\n",
    "assert any('Dropped 12 unused rows' in str(w.message) for w in issued_warnings)\n",
    "# models require futr_df and not provided raises an error\n",
    "test_fail(lambda: nf.predict(), contains=\"Models require the following future exogenous features: {'trend'}\") \n",
    "# missing feature in futr_df raises an error\n",
    "test_fail(lambda: nf.predict(futr_df=AirPassengersPanel_test.drop(columns='trend')), contains=\"missing from `futr_df`: {'trend'}\")\n",
    "# null values in futr_df raises an error\n",
    "test_fail(lambda: nf.predict(futr_df=AirPassengersPanel_test.assign(trend=np.nan)), contains='Found null values in `futr_df`')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e78b113",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Test inplace model fitting\n",
    "models = [MLP(h=12, input_size=12, max_steps=1, scaler_type='robust')]\n",
    "initial_weights = models[0].mlp[0].weight.detach().clone()\n",
    "fcst = NeuralForecast(models=models, freq='M')\n",
    "fcst.fit(df=AirPassengersPanel_train, static_df=AirPassengersStatic, use_init_models=True)\n",
    "after_weights = fcst.models_init[0].mlp[0].weight.detach().clone()\n",
    "assert np.allclose(initial_weights, after_weights), 'init models should not be modified'\n",
    "assert len(fcst.models[0].train_trajectories)>0, 'models stored trajectories should not be empty'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "340dd8a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Test predict_insample\n",
    "test_size = 12\n",
    "n_series = 2\n",
    "h = 12\n",
    "\n",
    "def get_expected_size(df, h, test_size, step_size):\n",
    "    expected_size = 0\n",
    "    uids = df['unique_id'].unique()\n",
    "    for uid in uids:\n",
    "        input_len = len(df[df['unique_id'] == uid])\n",
    "        expected_size += ((input_len - test_size - h) / step_size + 1)*h\n",
    "    return expected_size\n",
    "        \n",
    "models = [\n",
    "    NHITS(h=h, input_size=24, loss=MQLoss(level=[80]), max_steps=1, alias='NHITS', scaler_type=None),\n",
    "    RNN(h=h, input_size=-1, loss=MAE(), max_steps=1, alias='RNN', scaler_type=None),\n",
    "    ]\n",
    "\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "cv = nf.cross_validation(df=AirPassengersPanel_train, static_df=AirPassengersStatic, val_size=0, test_size=test_size, n_windows=None)\n",
    "\n",
    "forecasts = nf.predict_insample(step_size=1)\n",
    "\n",
    "expected_size = get_expected_size(AirPassengersPanel_train, h, test_size, step_size=1)\n",
    "assert len(forecasts) == expected_size, f'Shape mismatch in predict_insample: {len(forecasts)=}, {expected_size=}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d996a0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Test predict_insample (different lengths)\n",
    "diff_len_df = generate_series(n_series=n_series, max_length=100)\n",
    "\n",
    "nf = NeuralForecast(models=models, freq='D')\n",
    "cv = nf.cross_validation(df=diff_len_df, val_size=0, test_size=test_size, n_windows=None)\n",
    "\n",
    "forecasts = nf.predict_insample(step_size=1)\n",
    "expected_size = get_expected_size(diff_len_df, h, test_size, step_size=1)\n",
    "assert len(forecasts) == expected_size, f'Shape mismatch in predict_insample: {len(forecasts)=}, {expected_size=}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23fb98f8-0e27-44b2-a8a9-3b551986d53f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide,\n",
    "# Test predict_insample step_size\n",
    "\n",
    "h = 12\n",
    "train_end = AirPassengersPanel_train['ds'].max()\n",
    "sizes = AirPassengersPanel_train['unique_id'].value_counts().to_numpy()\n",
    "for step_size, test_size in [(7, 0), (9, 0), (7, 5), (9, 5)]:\n",
    "    models = [NHITS(h=h, input_size=12, max_steps=1)]\n",
    "    nf = NeuralForecast(models=models, freq='M')\n",
    "    nf.fit(AirPassengersPanel_train)\n",
    "    # Note: only apply set_test_size() upon nf.fit(), otherwise it would have set the test_size = 0\n",
    "    nf.models[0].set_test_size(test_size)\n",
    "    \n",
    "    forecasts = nf.predict_insample(step_size=step_size)\n",
    "    last_cutoff = train_end - test_size * pd.offsets.MonthEnd() - h * pd.offsets.MonthEnd()\n",
    "    n_expected_cutoffs = (sizes[0] - test_size - nf.h + step_size) // step_size\n",
    "\n",
    "    # compare cutoff values\n",
    "    expected_cutoffs = np.flip(np.array([last_cutoff - step_size * i * pd.offsets.MonthEnd() for i in range(n_expected_cutoffs)]))\n",
    "    actual_cutoffs = np.array([pd.Timestamp(x) for x in forecasts[forecasts['unique_id']==nf.uids[1]]['cutoff'].unique()])\n",
    "    np.testing.assert_array_equal(expected_cutoffs, actual_cutoffs, err_msg=f\"{step_size=},{expected_cutoffs=},{actual_cutoffs=}\")\n",
    "    \n",
    "    # check forecast-points count per series\n",
    "    cutoffs_by_series = forecasts.groupby(['unique_id', 'cutoff']).size().unstack('unique_id')\n",
    "    pd.testing.assert_series_equal(cutoffs_by_series['Airline1'], cutoffs_by_series['Airline2'], check_names=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "263c9c1bf36674b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide,\n",
    "# Test predict_insample\n",
    "\n",
    "def get_expected_cols(model, level):\n",
    "    # index columns\n",
    "    n_cols = 4\n",
    "    for model in models:\n",
    "        if isinstance(loss, (DistributionLoss, PMM, GMM, NBMM)):\n",
    "            if level is None:\n",
    "                # Variations of DistributionLoss return the sample mean as well\n",
    "                n_cols += len(loss.quantiles) + 1\n",
    "            else:\n",
    "                # Variations of DistributionLoss return the sample mean as well\n",
    "                n_cols += 2 * len(level) + 1\n",
    "        else:\n",
    "            if level is None:\n",
    "                # Other probabilistic models return the sample mean as well\n",
    "                n_cols += 1\n",
    "            # Other probabilistic models return just the levels\n",
    "            else:\n",
    "                n_cols += len(level) + 1\n",
    "    return n_cols\n",
    "\n",
    "for loss in [\n",
    "    # IQLoss(), \n",
    "    DistributionLoss(distribution=\"Normal\", level=[80]),\n",
    "    PMM(level=[80]),\n",
    "]:\n",
    "    for level in [None, [80, 90]]:\n",
    "        models = [\n",
    "            NHITS(h=12, input_size=12, loss=loss, max_steps=1),\n",
    "            LSTM(h=12, input_size=12, loss=loss, max_steps=1, recurrent=True),\n",
    "        ]\n",
    "        nf = NeuralForecast(models=models, freq='D')\n",
    "\n",
    "        nf.fit(df=AirPassengersPanel_train)\n",
    "        df = nf.predict_insample(step_size=1, level=level)\n",
    "        expected_cols = get_expected_cols(models, level)\n",
    "        assert df.shape[1] == expected_cols, f'Shape mismatch for {loss} and level={level} in predict_insample: cols={df.shape[1]}, expected_cols={expected_cols}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d6bbc2c-d38f-4cec-a3ef-15164852479f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Test aliases\n",
    "config_drnn = {'input_size': tune.choice([-1]), \n",
    "               'encoder_hidden_size': tune.choice([5, 10]),\n",
    "               'max_steps': 1,\n",
    "               'val_check_steps': 1,\n",
    "               'step_size': 1}\n",
    "models = [\n",
    "    # test Auto\n",
    "    AutoDilatedRNN(h=12, config=config_drnn, cpus=1, num_samples=2, alias='AutoDIL'),\n",
    "    # test BaseWindows\n",
    "    NHITS(h=12, input_size=24, loss=MQLoss(level=[80]), max_steps=1, alias='NHITSMQ'),\n",
    "    # test BaseRecurrent\n",
    "    RNN(h=12, input_size=-1, encoder_hidden_size=10, max_steps=1,\n",
    "            stat_exog_list=['airline1'],\n",
    "            futr_exog_list=['trend'], hist_exog_list=['y_[lag12]'], alias='MyRNN'),\n",
    "    # test BaseMultivariate\n",
    "    StemGNN(h=12, input_size=24, n_series=2, max_steps=1, scaler_type='robust', alias='StemMulti'),\n",
    "    # test model without alias\n",
    "    NHITS(h=12, input_size=24, max_steps=1),\n",
    "]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(df=AirPassengersPanel_train, static_df=AirPassengersStatic)\n",
    "forecasts = nf.predict(futr_df=AirPassengersPanel_test)\n",
    "test_eq(\n",
    "    forecasts.columns.to_list(),\n",
    "    ['unique_id', 'ds', 'AutoDIL', 'NHITSMQ-median', 'NHITSMQ-lo-80', 'NHITSMQ-hi-80', 'MyRNN', 'StemMulti', 'NHITS']\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d3779a6-2d03-4ac3-9f01-8bd5cb306845",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Unit test for core/model interactions\n",
    "config = {'input_size': tune.choice([12, 24]), \n",
    "          'hidden_size': 256,\n",
    "          'max_steps': 1,\n",
    "          'val_check_steps': 1,\n",
    "          'step_size': 12}\n",
    "\n",
    "config_drnn = {'input_size': tune.choice([-1]), \n",
    "               'encoder_hidden_size': tune.choice([5, 10]),\n",
    "               'max_steps': 1,\n",
    "               'val_check_steps': 1,\n",
    "               'step_size': 1}\n",
    "\n",
    "fcst = NeuralForecast(\n",
    "    models=[\n",
    "        AutoDilatedRNN(h=12, config=config_drnn, cpus=1, num_samples=2),\n",
    "        DeepAR(h=12, input_size=24, max_steps=1,\n",
    "               stat_exog_list=['airline1'], futr_exog_list=['trend']),\n",
    "        DilatedRNN(h=12, input_size=-1, encoder_hidden_size=10, max_steps=1,\n",
    "                   stat_exog_list=['airline1'],\n",
    "                   futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "        RNN(h=12, input_size=-1, encoder_hidden_size=10, max_steps=1,\n",
    "            inference_input_size=24,\n",
    "            stat_exog_list=['airline1'],\n",
    "            futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "        TCN(h=12, input_size=-1, encoder_hidden_size=10, max_steps=1,\n",
    "            stat_exog_list=['airline1'],\n",
    "            futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "        AutoMLP(h=12, config=config, cpus=1, num_samples=2),\n",
    "        NBEATSx(h=12, input_size=12, max_steps=1,\n",
    "                stat_exog_list=['airline1'],\n",
    "                futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "        NHITS(h=12, input_size=24, loss=MQLoss(level=[80]), max_steps=1),\n",
    "        NHITS(h=12, input_size=12, max_steps=1,\n",
    "              stat_exog_list=['airline1'],\n",
    "              futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "        DLinear(h=12, input_size=24, max_steps=1),\n",
    "        MLP(h=12, input_size=12, max_steps=1,\n",
    "            stat_exog_list=['airline1'],\n",
    "            futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "        TFT(h=12, input_size=24, max_steps=1),\n",
    "        VanillaTransformer(h=12, input_size=24, max_steps=1),\n",
    "        Informer(h=12, input_size=24, max_steps=1),\n",
    "        Autoformer(h=12, input_size=24, max_steps=1),\n",
    "        FEDformer(h=12, input_size=24, max_steps=1),\n",
    "        PatchTST(h=12, input_size=24, max_steps=1),\n",
    "        TimesNet(h=12, input_size=24, max_steps=1),\n",
    "        StemGNN(h=12, input_size=24, n_series=2, max_steps=1, scaler_type='robust'),\n",
    "        TSMixer(h=12, input_size=24, n_series=2, max_steps=1, scaler_type='robust'),\n",
    "        TSMixerx(h=12, input_size=24, n_series=2, max_steps=1, scaler_type='robust'),\n",
    "    ],\n",
    "    freq='M'\n",
    ")\n",
    "fcst.fit(df=AirPassengersPanel_train, static_df=AirPassengersStatic)\n",
    "forecasts = fcst.predict(futr_df=AirPassengersPanel_test)\n",
    "forecasts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40038532-fd68-4375-b7da-ba5bc2491c5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "fig, ax = plt.subplots(1, 1, figsize = (20, 7))\n",
    "plot_df = pd.concat([AirPassengersPanel_train, forecasts.reset_index()]).set_index('ds')\n",
    "\n",
    "plot_df[plot_df['unique_id']=='Airline1'].drop(['unique_id','trend','y_[lag12]'], axis=1).plot(ax=ax, linewidth=2)\n",
    "\n",
    "ax.set_title('AirPassengers Forecast', fontsize=22)\n",
    "ax.set_ylabel('Monthly Passengers', fontsize=20)\n",
    "ax.set_xlabel('Timestamp [t]', fontsize=20)\n",
    "ax.legend(prop={'size': 15})\n",
    "ax.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d61909b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "fig, ax = plt.subplots(1, 1, figsize = (20, 7))\n",
    "plot_df = pd.concat([AirPassengersPanel_train, forecasts.reset_index()]).set_index('ds')\n",
    "\n",
    "plot_df[plot_df['unique_id']=='Airline2'].drop(['unique_id','trend','y_[lag12]'], axis=1).plot(ax=ax, linewidth=2)\n",
    "\n",
    "ax.set_title('AirPassengers Forecast', fontsize=22)\n",
    "ax.set_ylabel('Monthly Passengers', fontsize=20)\n",
    "ax.set_xlabel('Timestamp [t]', fontsize=20)\n",
    "ax.legend(prop={'size': 15})\n",
    "ax.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe89e9e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def config_optuna(trial):\n",
    "    return {\"input_size\": trial.suggest_categorical('input_size', [12, 24]),\n",
    "        \"hist_exog_list\": trial.suggest_categorical('hist_exog_list', [['trend'], ['y_[lag12]'], ['trend', 'y_[lag12]']]),\n",
    "        \"futr_exog_list\": ['trend'],\n",
    "        \"max_steps\": 10,\n",
    "        \"val_check_steps\": 5}\n",
    "\n",
    "config_ray = {'input_size': tune.choice([12, 24]), \n",
    "          'hist_exog_list': tune.choice([['trend'], ['y_[lag12]'], ['trend', 'y_[lag12]']]),\n",
    "          'futr_exog_list': ['trend'],\n",
    "          'max_steps': 10,\n",
    "          'val_check_steps': 5}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85180620",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test training with an iterative dataset produces the same results as directly passing in the dataset as a pandas dataframe\n",
    "AirPassengersPanel_train['id'] = AirPassengersPanel_train['unique_id']\n",
    "AirPassengersPanel_test['id'] = AirPassengersPanel_test['unique_id']\n",
    "\n",
    "models = [\n",
    "    NHITS(h=12, input_size=12, max_steps=10, futr_exog_list=['trend'], random_seed=1),\n",
    "    AutoMLP(h=12, config=config_optuna, num_samples=2, backend='optuna', search_alg=optuna.samplers.TPESampler(seed=0)), # type: ignore\n",
    "    AutoNBEATSx(h=12, config=config_ray, cpus=1, num_samples=2)\n",
    "]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "\n",
    "# fit+predict with pandas dataframe\n",
    "nf.fit(df=AirPassengersPanel_train.drop(columns='unique_id'), use_init_models=True, id_col='id')\n",
    "pred_dataframe = nf.predict(futr_df=AirPassengersPanel_test.drop(columns='unique_id')).reset_index()\n",
    "\n",
    "# fit+predict with data directory\n",
    "with tempfile.TemporaryDirectory() as tmpdir:\n",
    "    AirPassengersPanel_train.to_parquet(tmpdir, partition_cols=['unique_id'], index=False)\n",
    "    data_directory = sorted([str(path) for path in Path(tmpdir).iterdir()])\n",
    "    nf.fit(df=data_directory, use_init_models=True, id_col='id')\n",
    "\n",
    "pred_df = AirPassengersPanel_train[AirPassengersPanel_train['unique_id'] == 'Airline2'].drop(columns='unique_id')\n",
    "futr_df = AirPassengersPanel_test[AirPassengersPanel_test['unique_id'] == 'Airline2'].drop(columns='unique_id')\n",
    "\n",
    "pred_iterative = nf.predict(df=pred_df, futr_df=futr_df)\n",
    "pred_airline2 = pred_dataframe[pred_dataframe['id'] == 'Airline2']\n",
    "np.testing.assert_allclose(pred_iterative['NHITS'], pred_airline2['NHITS'], rtol=0, atol=1)\n",
    "np.testing.assert_allclose(pred_iterative['AutoMLP'], pred_airline2['AutoMLP'], rtol=0, atol=1)\n",
    "np.testing.assert_allclose(pred_iterative['AutoNBEATSx'], pred_airline2['AutoNBEATSx'], rtol=0, atol=1)\n",
    "\n",
    "# remove id columns to not impact future tests\n",
    "AirPassengersPanel_train = AirPassengersPanel_train.drop(columns='id')\n",
    "AirPassengersPanel_test = AirPassengersPanel_test.drop(columns='id')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4a8162a-3d9d-48df-a314-3a2ce0377e36",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "config = {'input_size': tune.choice([12, 24]), \n",
    "          'hidden_size': 256,\n",
    "          'max_steps': 1,\n",
    "          'val_check_steps': 1,\n",
    "          'step_size': 12}\n",
    "\n",
    "config_drnn = {'input_size': tune.choice([-1]), \n",
    "               'encoder_hidden_size': tune.choice([5, 10]),\n",
    "               'max_steps': 1,\n",
    "               'val_check_steps': 1,\n",
    "               'step_size': 1}\n",
    "\n",
    "fcst = NeuralForecast(\n",
    "    models=[\n",
    "        DilatedRNN(h=12, input_size=-1, encoder_hidden_size=10, max_steps=1),\n",
    "        AutoMLP(h=12, config=config, cpus=1, num_samples=1),\n",
    "        NHITS(h=12, input_size=12, max_steps=1)\n",
    "    ],\n",
    "    freq='M'\n",
    ")\n",
    "cv_df = fcst.cross_validation(df=AirPassengersPanel, static_df=AirPassengersStatic, n_windows=3, step_size=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17c5ea12-ed87-4e46-ad04-3088e7167dfd",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "#test cross validation no leakage\n",
    "def test_cross_validation(df, static_df, h, test_size):\n",
    "    if (test_size - h) % 1:\n",
    "        raise Exception(\"`test_size - h` should be module `step_size`\")\n",
    "    \n",
    "    n_windows = int((test_size - h) / 1) + 1\n",
    "    Y_test_df = df.groupby('unique_id').tail(test_size)\n",
    "    Y_train_df = df.drop(Y_test_df.index)\n",
    "    config = {'input_size': tune.choice([12, 24]),\n",
    "              'step_size': 12, 'hidden_size': 256, 'max_steps': 1, 'val_check_steps': 1}\n",
    "    config_drnn = {'input_size': tune.choice([-1]), 'encoder_hidden_size': tune.choice([5, 10]),\n",
    "                   'max_steps': 1, 'val_check_steps': 1}\n",
    "    fcst = NeuralForecast(\n",
    "        models=[\n",
    "            AutoDilatedRNN(h=12, config=config_drnn, cpus=1, num_samples=1),\n",
    "            DilatedRNN(h=12, input_size=-1, encoder_hidden_size=5, max_steps=1),\n",
    "            RNN(h=12, input_size=-1, encoder_hidden_size=5, max_steps=1,\n",
    "                stat_exog_list=['airline1'], futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "            TCN(h=12, input_size=-1, encoder_hidden_size=5, max_steps=1,\n",
    "                stat_exog_list=['airline1'], futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "            AutoMLP(h=12, config=config, cpus=1, num_samples=1),\n",
    "            MLP(h=12, input_size=12, max_steps=1, scaler_type='robust'),\n",
    "            NBEATSx(h=12, input_size=12, max_steps=1,\n",
    "                    stat_exog_list=['airline1'], futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "            NHITS(h=12, input_size=12, max_steps=1, scaler_type='robust'),\n",
    "            NHITS(h=12, input_size=12, loss=MQLoss(level=[80]), max_steps=1),\n",
    "            TFT(h=12, input_size=24, max_steps=1, scaler_type='robust'),\n",
    "            DLinear(h=12, input_size=24, max_steps=1),\n",
    "            VanillaTransformer(h=12, input_size=12, max_steps=1, scaler_type=None),\n",
    "            Informer(h=12, input_size=12, max_steps=1, scaler_type=None),\n",
    "            Autoformer(h=12, input_size=12, max_steps=1, scaler_type=None),\n",
    "            FEDformer(h=12, input_size=12, max_steps=1, scaler_type=None),\n",
    "            PatchTST(h=12, input_size=24, max_steps=1, scaler_type=None),\n",
    "            TimesNet(h=12, input_size=24, max_steps=1, scaler_type='standard'),\n",
    "            StemGNN(h=12, input_size=12, n_series=2, max_steps=1, scaler_type='robust'),\n",
    "            TSMixer(h=12, input_size=12, n_series=2, max_steps=1, scaler_type='robust'),\n",
    "            TSMixerx(h=12, input_size=12, n_series=2, max_steps=1, scaler_type='robust'),\n",
    "            DeepAR(h=12, input_size=24, max_steps=1,\n",
    "               stat_exog_list=['airline1'], futr_exog_list=['trend']),\n",
    "        ],\n",
    "        freq='M'\n",
    "    )\n",
    "    fcst.fit(df=Y_train_df, static_df=static_df)\n",
    "    Y_hat_df = fcst.predict(futr_df=Y_test_df)\n",
    "    Y_hat_df = Y_hat_df.merge(Y_test_df, how='left', on=['unique_id', 'ds'])\n",
    "    last_dates = Y_train_df.groupby('unique_id').tail(1)\n",
    "    last_dates = last_dates[['unique_id', 'ds']].rename(columns={'ds': 'cutoff'})\n",
    "    Y_hat_df = Y_hat_df.merge(last_dates, how='left', on='unique_id')\n",
    "    \n",
    "    #cross validation\n",
    "    fcst = NeuralForecast(\n",
    "        models=[\n",
    "            AutoDilatedRNN(h=12, config=config_drnn, cpus=1, num_samples=1),\n",
    "            DilatedRNN(h=12, input_size=-1, encoder_hidden_size=5, max_steps=1),\n",
    "            RNN(h=12, input_size=-1, encoder_hidden_size=5, max_steps=1,\n",
    "                stat_exog_list=['airline1'], futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "            TCN(h=12, input_size=-1, encoder_hidden_size=5, max_steps=1,\n",
    "                stat_exog_list=['airline1'], futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "            AutoMLP(h=12, config=config, cpus=1, num_samples=1),\n",
    "            MLP(h=12, input_size=12, max_steps=1, scaler_type='robust'),\n",
    "            NBEATSx(h=12, input_size=12, max_steps=1,\n",
    "                    stat_exog_list=['airline1'], futr_exog_list=['trend'], hist_exog_list=['y_[lag12]']),\n",
    "            NHITS(h=12, input_size=12, max_steps=1, scaler_type='robust'),\n",
    "            NHITS(h=12, input_size=12, loss=MQLoss(level=[80]), max_steps=1),\n",
    "            TFT(h=12, input_size=24, max_steps=1, scaler_type='robust'),\n",
    "            DLinear(h=12, input_size=24, max_steps=1),\n",
    "            VanillaTransformer(h=12, input_size=12, max_steps=1, scaler_type=None),\n",
    "            Informer(h=12, input_size=12, max_steps=1, scaler_type=None),\n",
    "            Autoformer(h=12, input_size=12, max_steps=1, scaler_type=None),\n",
    "            FEDformer(h=12, input_size=12, max_steps=1, scaler_type=None),\n",
    "            PatchTST(h=12, input_size=24, max_steps=1, scaler_type=None),\n",
    "            TimesNet(h=12, input_size=24, max_steps=1, scaler_type='standard'),\n",
    "            StemGNN(h=12, input_size=12, n_series=2, max_steps=1, scaler_type='robust'),\n",
    "            TSMixer(h=12, input_size=12, n_series=2, max_steps=1, scaler_type='robust'),\n",
    "            TSMixerx(h=12, input_size=12, n_series=2, max_steps=1, scaler_type='robust'),\n",
    "            DeepAR(h=12, input_size=24, max_steps=1,\n",
    "               stat_exog_list=['airline1'], futr_exog_list=['trend']),\n",
    "        ],\n",
    "        freq='M'\n",
    "    )\n",
    "    Y_hat_df_cv = fcst.cross_validation(df, static_df=static_df, test_size=test_size, \n",
    "                                        n_windows=None)\n",
    "    for col in ['ds', 'cutoff']:\n",
    "        Y_hat_df_cv[col] = pd.to_datetime(Y_hat_df_cv[col].astype(str))\n",
    "        Y_hat_df[col] = pd.to_datetime(Y_hat_df[col].astype(str))\n",
    "    pd.testing.assert_frame_equal(\n",
    "        Y_hat_df[Y_hat_df_cv.columns],\n",
    "        Y_hat_df_cv,\n",
    "        check_dtype=False,\n",
    "        atol=1e-5,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0467904-748e-42ec-99cc-bac514626304",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "test_cross_validation(AirPassengersPanel, AirPassengersStatic, h=12, test_size=12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03396c73",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test cv with series of different sizes\n",
    "series = pd.DataFrame({\n",
    "    'unique_id': np.repeat([0, 1], [10, 15]),\n",
    "    'ds': np.arange(25),\n",
    "    'y': np.random.rand(25),\n",
    "})\n",
    "nf = NeuralForecast(\n",
    "    freq=1,\n",
    "    models=[MLP(input_size=5, h=5, max_steps=0, enable_progress_bar=False)]\n",
    ")\n",
    "cv_df = nf.cross_validation(df=series, n_windows=3, step_size=5)\n",
    "expected = pd.DataFrame({\n",
    "    'unique_id': np.repeat([0, 1], [5, 10]),\n",
    "    'ds': np.hstack([np.arange(5, 10), np.arange(15, 25)]),\n",
    "    'cutoff': np.repeat([4, 14, 19], 5)\n",
    "})\n",
    "expected = expected.merge(series, on=['unique_id', 'ds'])\n",
    "pd.testing.assert_frame_equal(expected, cv_df.drop(columns='MLP'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f61d030-a51e-49f8-a16e-b97bccb62401",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test save and load\n",
    "config = {'input_size': tune.choice([12, 24]),\n",
    "          'hidden_size': 256,\n",
    "          'max_steps': 1,\n",
    "          'val_check_steps': 1,\n",
    "          'step_size': 12}\n",
    "\n",
    "config_drnn = {'input_size': tune.choice([-1]),\n",
    "               'encoder_hidden_size': tune.choice([5, 10]),\n",
    "               'max_steps': 1,\n",
    "               'val_check_steps': 1}\n",
    "\n",
    "fcst = NeuralForecast(\n",
    "    models=[\n",
    "        AutoRNN(h=12, config=config_drnn, cpus=1, num_samples=2, refit_with_val=True),\n",
    "        DilatedRNN(h=12, input_size=-1, encoder_hidden_size=5, max_steps=1),\n",
    "        AutoMLP(h=12, config=config, cpus=1, num_samples=2),\n",
    "        NHITS(h=12, input_size=12, max_steps=1,\n",
    "              futr_exog_list=['trend'], hist_exog_list=['y_[lag12]'], alias='Model1'),\n",
    "        StemGNN(h=12, input_size=12, n_series=2, max_steps=1, scaler_type='robust')\n",
    "    ],\n",
    "    freq='M'\n",
    ")\n",
    "prediction_intervals = PredictionIntervals()\n",
    "fcst.fit(AirPassengersPanel_train, prediction_intervals=prediction_intervals)\n",
    "forecasts1 = fcst.predict(futr_df=AirPassengersPanel_test, level=[50])\n",
    "save_paths = ['./examples/debug_run/']\n",
    "try:\n",
    "    s3fs.S3FileSystem().ls('s3://nixtla-tmp')    \n",
    "    pyver = f'{sys.version_info.major}_{sys.version_info.minor}'\n",
    "    sha = git.Repo(search_parent_directories=True).head.object.hexsha\n",
    "    save_dir = f'{sys.platform}-{pyver}-{sha}'\n",
    "    save_paths.append(f's3://nixtla-tmp/neural/{save_dir}')\n",
    "except Exception as e:\n",
    "    print(e)\n",
    "\n",
    "for path in save_paths:\n",
    "    fcst.save(path=path, model_index=None, overwrite=True, save_dataset=True)\n",
    "    fcst2 = NeuralForecast.load(path=path)\n",
    "    forecasts2 = fcst2.predict(futr_df=AirPassengersPanel_test, level=[50])\n",
    "    pd.testing.assert_frame_equal(forecasts1, forecasts2[forecasts1.columns])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d221d90-7a89-4338-96b8-09543f3d5554",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test save and load without dataset\n",
    "shutil.rmtree('examples/debug_run')\n",
    "fcst = NeuralForecast(\n",
    "    models=[DilatedRNN(h=12, input_size=-1, encoder_hidden_size=5, max_steps=1)],\n",
    "    freq='M',\n",
    ")\n",
    "fcst.fit(AirPassengersPanel_train)\n",
    "forecasts1 = fcst.predict(futr_df=AirPassengersPanel_test)\n",
    "fcst.save(path='./examples/debug_run/', model_index=None, overwrite=True, save_dataset=False)\n",
    "fcst2 = NeuralForecast.load(path='./examples/debug_run/')\n",
    "forecasts2 = fcst2.predict(df=AirPassengersPanel_train, futr_df=AirPassengersPanel_test)\n",
    "np.testing.assert_allclose(forecasts1['DilatedRNN'], forecasts2['DilatedRNN'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c22ad495",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test `enable_checkpointing=True` should generate chkpt\n",
    "shutil.rmtree('lightning_logs')\n",
    "fcst = NeuralForecast(\n",
    "    models=[\n",
    "        MLP(h=12, input_size=12, max_steps=10, val_check_steps=5, enable_checkpointing=True),\n",
    "        RNN(h=12, input_size=-1, max_steps=10, val_check_steps=5, enable_checkpointing=True)\n",
    "    ],\n",
    "    freq='M'\n",
    ")\n",
    "fcst.fit(AirPassengersPanel_train)\n",
    "last_log = f\"lightning_logs/{os.listdir('lightning_logs')[-1]}\"\n",
    "no_chkpt_found = ~np.any([file.endswith('checkpoints') for file in os.listdir(last_log)])\n",
    "test_eq(no_chkpt_found, False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ac7a0b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test `enable_checkpointing=False` should not generate chkpt\n",
    "shutil.rmtree('lightning_logs')\n",
    "fcst = NeuralForecast(\n",
    "    models=[\n",
    "        MLP(h=12, input_size=12, max_steps=10, val_check_steps=5),\n",
    "        RNN(h=12, input_size=-1, max_steps=10, val_check_steps=5)\n",
    "    ],\n",
    "    freq='M'\n",
    ")\n",
    "fcst.fit(AirPassengersPanel_train)\n",
    "last_log = f\"lightning_logs/{os.listdir('lightning_logs')[-1]}\"\n",
    "no_chkpt_found = ~np.any([file.endswith('checkpoints') for file in os.listdir(last_log)])\n",
    "test_eq(no_chkpt_found, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8602941e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test short time series\n",
    "config = {'input_size': tune.choice([12, 24]), \n",
    "          'max_steps': 1,\n",
    "          'val_check_steps': 1}\n",
    "\n",
    "fcst = NeuralForecast(\n",
    "    models=[\n",
    "        AutoNBEATS(h=12, config=config, cpus=1, num_samples=2)],\n",
    "    freq='M'\n",
    ")\n",
    "\n",
    "AirPassengersShort = AirPassengersPanel.tail(36+144).reset_index(drop=True)\n",
    "forecasts = fcst.cross_validation(AirPassengersShort, val_size=48, n_windows=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cadac88d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test validation scale BaseWindows\n",
    "\n",
    "models = [NHITS(h=12, input_size=24, max_steps=50, scaler_type='robust')]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(AirPassengersPanel_train,val_size=12)\n",
    "valid_losses = nf.models[0].valid_trajectories\n",
    "assert valid_losses[-1][1] < 40, 'Validation loss is too high'\n",
    "assert valid_losses[-1][1] > 10, 'Validation loss is too low'\n",
    "\n",
    "models = [NHITS(h=12, input_size=24, max_steps=50, scaler_type=None)]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(AirPassengersPanel_train,val_size=12)\n",
    "valid_losses = nf.models[0].valid_trajectories\n",
    "assert valid_losses[-1][1] < 40, 'Validation loss is too high'\n",
    "assert valid_losses[-1][1] > 10, 'Validation loss is too low'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee083d85",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test validation scale BaseRecurrent\n",
    "\n",
    "nf = NeuralForecast(\n",
    "    models=[LSTM(h=12,\n",
    "                 input_size=-1,\n",
    "                 loss=MAE(),\n",
    "                 scaler_type='robust',\n",
    "                 encoder_n_layers=2,\n",
    "                 encoder_hidden_size=128,\n",
    "                 context_size=10,\n",
    "                 decoder_hidden_size=128,\n",
    "                 decoder_layers=2,\n",
    "                 max_steps=50,\n",
    "                 val_check_steps=10,\n",
    "                 )\n",
    "    ],\n",
    "    freq='M'\n",
    ")\n",
    "nf.fit(AirPassengersPanel_train,val_size=12)\n",
    "valid_losses = nf.models[0].valid_trajectories\n",
    "assert valid_losses[-1][1] < 100, 'Validation loss is too high'\n",
    "assert valid_losses[-1][1] > 30, 'Validation loss is too low'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02ee53b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Test order of variables does not affect validation loss\n",
    "\n",
    "AirPassengersPanel_train['zeros'] = 0\n",
    "AirPassengersPanel_train['large_number'] = 100000\n",
    "AirPassengersPanel_train['available_mask'] = 1\n",
    "AirPassengersPanel_train = AirPassengersPanel_train[['unique_id','ds','zeros','y','available_mask','large_number']]\n",
    "\n",
    "models = [NHITS(h=12, input_size=24, max_steps=50, scaler_type='robust')]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(AirPassengersPanel_train,val_size=12)\n",
    "valid_losses = nf.models[0].valid_trajectories\n",
    "assert valid_losses[-1][1] < 40, 'Validation loss is too high'\n",
    "assert valid_losses[-1][1] > 10, 'Validation loss is too low'\n",
    "\n",
    "models = [NHITS(h=12, input_size=24, max_steps=50, scaler_type=None)]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(AirPassengersPanel_train,val_size=12)\n",
    "valid_losses = nf.models[0].valid_trajectories\n",
    "assert valid_losses[-1][1] < 40, 'Validation loss is too high'\n",
    "assert valid_losses[-1][1] > 10, 'Validation loss is too low'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ba31378",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Test fit fails if variable not in dataframe\n",
    "\n",
    "# Base Windows\n",
    "models = [NHITS(h=12, input_size=24, max_steps=1, hist_exog_list=['not_included'], scaler_type='robust')]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "test_fail(nf.fit,\n",
    "          contains='historical exogenous variables not found in input dataset',\n",
    "          args=(AirPassengersPanel_train,))\n",
    "\n",
    "models = [NHITS(h=12, input_size=24, max_steps=1, futr_exog_list=['not_included'], scaler_type='robust')]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "test_fail(nf.fit,\n",
    "          contains='future exogenous variables not found in input dataset',\n",
    "          args=(AirPassengersPanel_train,))\n",
    "\n",
    "models = [NHITS(h=12, input_size=24, max_steps=1, stat_exog_list=['not_included'], scaler_type='robust')]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "test_fail(nf.fit,\n",
    "          contains='static exogenous variables not found in input dataset',\n",
    "          args=(AirPassengersPanel_train,))\n",
    "\n",
    "# Base Recurrent\n",
    "models = [LSTM(h=12, input_size=24, max_steps=1, hist_exog_list=['not_included'], scaler_type='robust')]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "test_fail(nf.fit,\n",
    "          contains='historical exogenous variables not found in input dataset',\n",
    "          args=(AirPassengersPanel_train,))\n",
    "\n",
    "models = [LSTM(h=12, input_size=24, max_steps=1, futr_exog_list=['not_included'], scaler_type='robust')]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "test_fail(nf.fit,\n",
    "          contains='future exogenous variables not found in input dataset',\n",
    "          args=(AirPassengersPanel_train,))\n",
    "\n",
    "models = [LSTM(h=12, input_size=24, max_steps=1, stat_exog_list=['not_included'], scaler_type='robust')]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "test_fail(nf.fit,\n",
    "          contains='static exogenous variables not found in input dataset',\n",
    "          args=(AirPassengersPanel_train,))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d221479c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Test passing unused variables in dataframe does not affect forecasts  \n",
    "\n",
    "models = [NHITS(h=12, input_size=24, max_steps=5, hist_exog_list=['zeros'], scaler_type='robust')]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(AirPassengersPanel_train)\n",
    "\n",
    "Y_hat1 = nf.predict(df=AirPassengersPanel_train[['unique_id','ds','y','zeros','large_number']])\n",
    "Y_hat2 = nf.predict(df=AirPassengersPanel_train[['unique_id','ds','y','zeros']])\n",
    "\n",
    "pd.testing.assert_frame_equal(\n",
    "    Y_hat1,\n",
    "    Y_hat2,\n",
    "    check_dtype=False,\n",
    ")\n",
    "\n",
    "models = [LSTM(h=12, input_size=24, max_steps=5, hist_exog_list=['zeros'], scaler_type='robust')]\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(AirPassengersPanel_train)\n",
    "\n",
    "Y_hat1 = nf.predict(df=AirPassengersPanel_train[['unique_id','ds','y','zeros','large_number']])\n",
    "Y_hat2 = nf.predict(df=AirPassengersPanel_train[['unique_id','ds','y','zeros']])\n",
    "\n",
    "pd.testing.assert_frame_equal(\n",
    "    Y_hat1,\n",
    "    Y_hat2,\n",
    "    check_dtype=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30890c07-1763-4795-afba-f5ed916245be",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "#| polars\n",
    "import polars\n",
    "from polars.testing import assert_frame_equal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad51c803",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "#| polars\n",
    "renamer = {'unique_id': 'uid', 'ds': 'time', 'y': 'target'}\n",
    "inverse_renamer = {v: k for k, v in renamer.items()}\n",
    "AirPassengers_pl = polars.from_pandas(AirPassengersPanel_train)\n",
    "AirPassengers_pl = AirPassengers_pl.rename(renamer)\n",
    "AirPassengersStatic_pl = polars.from_pandas(AirPassengersStatic)\n",
    "AirPassengersStatic_pl = AirPassengersStatic_pl.rename({'unique_id': 'uid'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7502d18c-62a5-4381-bfdb-bba5300c5290",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "#| polars\n",
    "models = [LSTM(h=12, input_size=24, max_steps=5, scaler_type='robust')]\n",
    "\n",
    "# Pandas\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(AirPassengersPanel_train, static_df=AirPassengersStatic)\n",
    "insample_preds = nf.predict_insample()\n",
    "preds = nf.predict()\n",
    "cv_res = nf.cross_validation(df=AirPassengersPanel_train, static_df=AirPassengersStatic)\n",
    "\n",
    "# Polars\n",
    "nf = NeuralForecast(models=models, freq='1mo')\n",
    "nf.fit(\n",
    "    AirPassengers_pl,\n",
    "    static_df=AirPassengersStatic_pl,\n",
    "    id_col='uid',\n",
    "    time_col='time',\n",
    "    target_col='target',\n",
    ")\n",
    "insample_preds_pl = nf.predict_insample()\n",
    "preds_pl = nf.predict()\n",
    "cv_res_pl = nf.cross_validation(\n",
    "    df=AirPassengers_pl,\n",
    "    static_df=AirPassengersStatic_pl,\n",
    "    id_col='uid',\n",
    "    time_col='time',\n",
    "    target_col='target',\n",
    ")\n",
    "\n",
    "def assert_equal_dfs(pandas_df, polars_df):\n",
    "    mapping = {k: v for k, v in inverse_renamer.items() if k in polars_df}\n",
    "    pd.testing.assert_frame_equal(\n",
    "        pandas_df,\n",
    "        polars_df.rename(mapping).to_pandas(),\n",
    "    )\n",
    "\n",
    "assert_equal_dfs(preds, preds_pl)\n",
    "assert_equal_dfs(insample_preds, insample_preds_pl)\n",
    "assert_equal_dfs(cv_res, cv_res_pl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "906fd509-82f5-431e-86df-d73e928ddeef",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide,\n",
    "#| polars\n",
    "# Test predict_insample step_size\n",
    "\n",
    "h = 12\n",
    "train_end = AirPassengers_pl['time'].max()\n",
    "sizes = AirPassengers_pl['uid'].value_counts().to_numpy()\n",
    "\n",
    "for step_size, test_size in [(7, 0), (9, 0), (7, 5), (9, 5)]:\n",
    "    models = [NHITS(h=h, input_size=12, max_steps=1)]\n",
    "    nf = NeuralForecast(models=models, freq='1mo')\n",
    "    nf.fit(\n",
    "        AirPassengers_pl,\n",
    "        id_col='uid',\n",
    "        time_col='time',\n",
    "        target_col='target',    \n",
    "    )\n",
    "    # Note: only apply set_test_size() upon nf.fit(), otherwise it would have set the test_size = 0\n",
    "    nf.models[0].set_test_size(test_size)    \n",
    "    \n",
    "    forecasts = nf.predict_insample(step_size=step_size)\n",
    "    n_expected_cutoffs = (sizes[0][1] - test_size - nf.h + step_size) // step_size\n",
    "\n",
    "    # compare cutoff values\n",
    "    last_cutoff = train_end - test_size * pd.offsets.MonthEnd() - h * pd.offsets.MonthEnd()\n",
    "    expected_cutoffs = np.flip(np.array([last_cutoff - step_size * i * pd.offsets.MonthEnd() for i in range(n_expected_cutoffs)]))\n",
    "    pl_cutoffs = forecasts.filter(polars.col('uid') ==nf.uids[1]).select('cutoff').unique(maintain_order=True)\n",
    "    actual_cutoffs = np.sort(np.array([pd.Timestamp(x['cutoff']) for x in pl_cutoffs.rows(named=True)]))\n",
    "    np.testing.assert_array_equal(expected_cutoffs, actual_cutoffs, err_msg=f\"{step_size=},{expected_cutoffs=},{actual_cutoffs=}\")\n",
    "\n",
    "    # check forecast-points count per series\n",
    "    cutoffs_by_series = forecasts.group_by(['uid', 'cutoff']).count()\n",
    "    assert_frame_equal(cutoffs_by_series.filter(polars.col('uid') == \"Airline1\").select(['cutoff', 'count']), cutoffs_by_series.filter(polars.col('uid') == \"Airline2\").select(['cutoff', 'count'] ), check_row_order=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fa887b3-4164-4758-931d-8d28a71b19b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Test if any of the inputs contains NaNs with available_mask = 1, fit shall raise error\n",
    "# input type is pandas.DataFrame\n",
    "# available_mask is explicitly given\n",
    "\n",
    "n_static_features = 2\n",
    "n_temporal_features = 4\n",
    "temporal_df, static_df = generate_series(n_series=4,\n",
    "                                         min_length=50,\n",
    "                                         max_length=50,\n",
    "                                         n_static_features=n_static_features,\n",
    "                                         n_temporal_features=n_temporal_features, \n",
    "                                         equal_ends=False) \n",
    "temporal_df[\"available_mask\"] = 1\n",
    "temporal_df.loc[10:20, \"available_mask\"] = 0\n",
    "models = [NHITS(h=12, input_size=24, max_steps=20)]\n",
    "nf = NeuralForecast(models=models, freq='D')\n",
    "\n",
    "# test case 1: target has NaN values\n",
    "test_df1 = temporal_df.copy()\n",
    "test_df1.loc[5:7, \"y\"] = np.nan\n",
    "test_fail(lambda: nf.fit(test_df1), contains=\"Found missing values in ['y']\")\n",
    "\n",
    "# test case 2: exogenous has NaN values that are correctly flagged with exception\n",
    "test_df2 = temporal_df.copy()\n",
    "# temporal_0 won't raise ValueError as available_mask = 0\n",
    "test_df2.loc[15:18, \"temporal_0\"] = np.nan\n",
    "test_df2.loc[5, \"temporal_1\"] = np.nan\n",
    "test_df2.loc[25, \"temporal_2\"] = np.nan\n",
    "test_fail(lambda: nf.fit(test_df2), contains=\"Found missing values in ['temporal_1', 'temporal_2']\")\n",
    "\n",
    "# test case 3: static column has NaN values\n",
    "test_df3 = static_df.copy()\n",
    "test_df3.loc[3, \"static_1\"] = np.nan\n",
    "test_fail(lambda: nf.fit(temporal_df, static_df=test_df3), contains=\"Found missing values in ['static_1']\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a157b6b4-0943-48f9-9427-fa8cf0b15d49",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "#| polars\n",
    "# Test if any of the inputs contains NaNs with available_mask = 1, fit shall raise error\n",
    "# input type is polars.Dataframe\n",
    "# Note that available_mask is not explicitly provided for this test\n",
    "\n",
    "pl_df = polars.DataFrame(\n",
    "    {\n",
    "        'unique_id': [1]*50,\n",
    "        'y': list(range(50)), \n",
    "        'temporal_0': list(range(100,150)),\n",
    "        'temporal_1': list(range(200,250)),\n",
    "        'ds': polars.date_range(start=date(2022, 1, 1), end=date(2022, 2, 19), interval=\"1d\", eager=True), \n",
    "    }\n",
    ")\n",
    "\n",
    "pl_static_df = polars.DataFrame(\n",
    "    {\n",
    "        'unique_id': [1],\n",
    "        'static_0': [1.2], \n",
    "        'static_1': [10.9],\n",
    "    }\n",
    ")\n",
    "\n",
    "models = [NHITS(h=12, input_size=24, max_steps=20)]\n",
    "nf = NeuralForecast(models=models, freq='1d')\n",
    "\n",
    "# test case 1: target has NaN values\n",
    "test_pl_df1 = pl_df.clone()\n",
    "test_pl_df1[3, 'y'] = np.nan\n",
    "test_pl_df1[4, 'y'] = None\n",
    "test_fail(lambda: nf.fit(test_pl_df1), contains=\"Found missing values in ['y']\")\n",
    "\n",
    "# test case 2: exogenous has NaN values that are correctly flagged with exception\n",
    "test_pl_df2 = pl_df.clone()\n",
    "test_pl_df2[15, \"temporal_0\"] = np.nan\n",
    "test_pl_df2[5, \"temporal_1\"] = np.nan\n",
    "test_fail(lambda: nf.fit(test_pl_df2), contains=\"Found missing values in ['temporal_0', 'temporal_1']\")\n",
    "\n",
    "# test case 3: static column has NaN values\n",
    "test_pl_df3 = pl_static_df.clone()\n",
    "test_pl_df3[0, \"static_1\"] = np.nan\n",
    "test_fail(lambda: nf.fit(pl_df, static_df=test_pl_df3), contains=\"Found missing values in ['static_1']\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "859a474c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test customized optimizer behavior such that the user defined optimizer result should differ from default\n",
    "# tests consider models implemented using different base classes such as BaseWindows, BaseRecurrent, BaseMultivariate\n",
    "\n",
    "for nf_model in [NHITS, RNN, StemGNN]:\n",
    "    # default optimizer is based on Adam\n",
    "    params = {\"h\": 12, \"input_size\": 24, \"max_steps\": 1}\n",
    "    if nf_model.__name__ == \"StemGNN\":\n",
    "        params.update({\"n_series\": 2})\n",
    "    models = [nf_model(**params)]\n",
    "    nf = NeuralForecast(models=models, freq='M')\n",
    "    nf.fit(AirPassengersPanel_train)\n",
    "    default_optimizer_predict = nf.predict()\n",
    "    mean = default_optimizer_predict.loc[:, nf_model.__name__].mean()\n",
    "\n",
    "    # using a customized optimizer\n",
    "    params.update({\n",
    "        \"optimizer\": torch.optim.Adadelta,\n",
    "        \"optimizer_kwargs\": {\"rho\": 0.45}, \n",
    "    })\n",
    "    models2 = [nf_model(**params)]\n",
    "    nf2 = NeuralForecast(models=models2, freq='M')\n",
    "    nf2.fit(AirPassengersPanel_train)\n",
    "    customized_optimizer_predict = nf2.predict()\n",
    "    mean2 = customized_optimizer_predict.loc[:, nf_model.__name__].mean()\n",
    "    assert mean2 != mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3db3fe1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test that if the user-defined optimizer is not a subclass of torch.optim.optimizer, failed with exception\n",
    "# tests cover different types of base classes such as BaseWindows, BaseRecurrent, BaseMultivariate\n",
    "test_fail(lambda: NHITS(h=12, input_size=24, max_steps=10, optimizer=torch.nn.Module), contains=\"optimizer is not a valid subclass of torch.optim.Optimizer\")\n",
    "test_fail(lambda: RNN(h=12, input_size=24, max_steps=10, optimizer=torch.nn.Module), contains=\"optimizer is not a valid subclass of torch.optim.Optimizer\")\n",
    "test_fail(lambda: StemGNN(h=12, input_size=24, max_steps=10, n_series=2, optimizer=torch.nn.Module), contains=\"optimizer is not a valid subclass of torch.optim.Optimizer\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d908240f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test that if we pass \"lr\" parameter, we expect warning and it ignores the passed in 'lr' parameter\n",
    "# tests consider models implemented using different base classes such as BaseWindows, BaseRecurrent, BaseMultivariate\n",
    "\n",
    "for nf_model in [NHITS, RNN, StemGNN]:\n",
    "    params = {\n",
    "        \"h\": 12, \n",
    "        \"input_size\": 24, \n",
    "        \"max_steps\": 1, \n",
    "        \"optimizer\": torch.optim.Adadelta, \n",
    "        \"optimizer_kwargs\": {\"lr\": 0.8, \"rho\": 0.45}\n",
    "    }\n",
    "    if nf_model.__name__ == \"StemGNN\":\n",
    "        params.update({\"n_series\": 2})\n",
    "    models = [nf_model(**params)]\n",
    "    nf = NeuralForecast(models=models, freq='M')\n",
    "    with warnings.catch_warnings(record=True) as issued_warnings:\n",
    "        warnings.simplefilter('always', UserWarning)\n",
    "        nf.fit(AirPassengersPanel_train)\n",
    "        assert any(\"ignoring learning rate passed in optimizer_kwargs, using the model's learning rate\" in str(w.message) for w in issued_warnings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c97858b5-e6a0-4353-a48f-5a5460eb2314",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test that if we pass \"optimizer_kwargs\" but not \"optimizer\", we expect a warning\n",
    "# tests consider models implemented using different base classes such as BaseWindows, BaseRecurrent, BaseMultivariate\n",
    "\n",
    "for nf_model in [NHITS, RNN, StemGNN]:\n",
    "    params = {\n",
    "        \"h\": 12, \n",
    "        \"input_size\": 24, \n",
    "        \"max_steps\": 1,\n",
    "        \"optimizer_kwargs\": {\"lr\": 0.8, \"rho\": 0.45}\n",
    "    }\n",
    "    if nf_model.__name__ == \"StemGNN\":\n",
    "        params.update({\"n_series\": 2})\n",
    "    models = [nf_model(**params)]\n",
    "    nf = NeuralForecast(models=models, freq='M')\n",
    "    with warnings.catch_warnings(record=True) as issued_warnings:\n",
    "        warnings.simplefilter('always', UserWarning)\n",
    "        nf.fit(AirPassengersPanel_train)\n",
    "        assert any(\"ignoring optimizer_kwargs as the optimizer is not specified\" in str(w.message) for w in issued_warnings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24142322",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test customized lr_scheduler behavior such that the user defined lr_scheduler result should differ from default\n",
    "# tests consider models implemented using different base classes such as BaseWindows, BaseRecurrent, BaseMultivariate\n",
    "\n",
    "for nf_model in [NHITS, RNN, StemGNN]:\n",
    "    params = {\"h\": 12, \"input_size\": 24, \"max_steps\": 1}\n",
    "    if nf_model.__name__ == \"StemGNN\":\n",
    "        params.update({\"n_series\": 2})\n",
    "    models = [nf_model(**params)]\n",
    "    nf = NeuralForecast(models=models, freq='M')\n",
    "    nf.fit(AirPassengersPanel_train)\n",
    "    default_optimizer_predict = nf.predict()\n",
    "    mean = default_optimizer_predict.loc[:, nf_model.__name__].mean()\n",
    "\n",
    "    # using a customized lr_scheduler, default is StepLR\n",
    "    params.update({\n",
    "        \"lr_scheduler\": torch.optim.lr_scheduler.ConstantLR,\n",
    "        \"lr_scheduler_kwargs\": {\"factor\": 0.78}, \n",
    "    })\n",
    "    models2 = [nf_model(**params)]\n",
    "    nf2 = NeuralForecast(models=models2, freq='M')\n",
    "    nf2.fit(AirPassengersPanel_train)\n",
    "    customized_optimizer_predict = nf2.predict()\n",
    "    mean2 = customized_optimizer_predict.loc[:, nf_model.__name__].mean()\n",
    "    assert mean2 != mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54c7b5e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test that if the user-defined lr_scheduler is not a subclass of torch.optim.lr_scheduler, failed with exception\n",
    "# tests cover different types of base classes such as BaseWindows, BaseRecurrent, BaseMultivariate\n",
    "test_fail(lambda: NHITS(h=12, input_size=24, max_steps=10, lr_scheduler=torch.nn.Module), contains=\"lr_scheduler is not a valid subclass of torch.optim.lr_scheduler.LRScheduler\")\n",
    "test_fail(lambda: RNN(h=12, input_size=24, max_steps=10, lr_scheduler=torch.nn.Module), contains=\"lr_scheduler is not a valid subclass of torch.optim.lr_scheduler.LRScheduler\")\n",
    "test_fail(lambda: StemGNN(h=12, input_size=24, max_steps=10, n_series=2, lr_scheduler=torch.nn.Module), contains=\"lr_scheduler is not a valid subclass of torch.optim.lr_scheduler.LRScheduler\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1d8bebb",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test that if we pass in \"optimizer\" parameter, we expect warning and it ignores them\n",
    "# tests consider models implemented using different base classes such as BaseWindows, BaseRecurrent, BaseMultivariate\n",
    "\n",
    "for nf_model in [NHITS, RNN, StemGNN]:\n",
    "    params = {\n",
    "        \"h\": 12, \n",
    "        \"input_size\": 24, \n",
    "        \"max_steps\": 1, \n",
    "        \"lr_scheduler\": torch.optim.lr_scheduler.ConstantLR, \n",
    "        \"lr_scheduler_kwargs\": {\"optimizer\": torch.optim.Adadelta, \"factor\": 0.22}\n",
    "    }\n",
    "    if nf_model.__name__ == \"StemGNN\":\n",
    "        params.update({\"n_series\": 2})\n",
    "    models = [nf_model(**params)]\n",
    "    nf = NeuralForecast(models=models, freq='M')\n",
    "    with warnings.catch_warnings(record=True) as issued_warnings:\n",
    "        warnings.simplefilter('always', UserWarning)\n",
    "        nf.fit(AirPassengersPanel_train)\n",
    "        assert any(\"ignoring optimizer passed in lr_scheduler_kwargs, using the model's optimizer\" in str(w.message) for w in issued_warnings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06febece",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test that if we pass in \"lr_scheduler_kwargs\" but not \"lr_scheduler\", we expect a warning\n",
    "# tests consider models implemented using different base classes such as BaseWindows, BaseRecurrent, BaseMultivariate\n",
    "\n",
    "for nf_model in [NHITS, RNN, StemGNN]:\n",
    "    params = {\n",
    "        \"h\": 12, \n",
    "        \"input_size\": 24, \n",
    "        \"max_steps\": 1,\n",
    "        \"lr_scheduler_kwargs\": {\"optimizer\": torch.optim.Adadelta, \"factor\": 0.22}\n",
    "    }\n",
    "    if nf_model.__name__ == \"StemGNN\":\n",
    "        params.update({\"n_series\": 2})\n",
    "    models = [nf_model(**params)]\n",
    "    nf = NeuralForecast(models=models, freq='M')\n",
    "    with warnings.catch_warnings(record=True) as issued_warnings:\n",
    "        warnings.simplefilter('always', UserWarning)\n",
    "        nf.fit(AirPassengersPanel_train)\n",
    "        assert any(\"ignoring lr_scheduler_kwargs as the lr_scheduler is not specified\" in str(w.message) for w in issued_warnings)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0441e9d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test conformal prediction, method=conformal_distribution\n",
    "\n",
    "prediction_intervals = PredictionIntervals()\n",
    "\n",
    "models = []\n",
    "for nf_model in [NHITS, RNN, TSMixer]:\n",
    "    params = {\"h\": 12, \"input_size\": 24, \"max_steps\": 1}\n",
    "    if nf_model.__name__ == \"TSMixer\":\n",
    "        params.update({\"n_series\": 2})\n",
    "    models.append(nf_model(**params))\n",
    "\n",
    "\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(AirPassengersPanel_train, prediction_intervals=prediction_intervals)\n",
    "preds = nf.predict(futr_df=AirPassengersPanel_test, level=[90])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9eb89f3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "#| polars\n",
    "# test conformal prediction works for polar dataframe\n",
    "\n",
    "prediction_intervals = PredictionIntervals()\n",
    "\n",
    "models = []\n",
    "for nf_model in [NHITS, RNN, TSMixer]:\n",
    "    params = {\"h\": 12, \"input_size\": 24, \"max_steps\": 1}\n",
    "    if nf_model.__name__ == \"TSMixer\":\n",
    "        params.update({\"n_series\": 2})\n",
    "    models.append(nf_model(**params))\n",
    "\n",
    "\n",
    "nf = NeuralForecast(models=models, freq='1mo')\n",
    "nf.fit(AirPassengers_pl, prediction_intervals=prediction_intervals, time_col='time', id_col='uid', target_col='target')\n",
    "preds = nf.predict(level=[90])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0db88cac",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test conformal prediction, method=conformal_error\n",
    "\n",
    "prediction_intervals = PredictionIntervals(method=\"conformal_error\")\n",
    "\n",
    "models = []\n",
    "for nf_model in [NHITS, RNN, TSMixer]:\n",
    "    params = {\"h\": 12, \"input_size\": 24, \"max_steps\": 1}\n",
    "    if nf_model.__name__ == \"TSMixer\":\n",
    "        params.update({\"n_series\": 2})\n",
    "    models.append(nf_model(**params))\n",
    "\n",
    "\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(AirPassengersPanel_train, prediction_intervals=prediction_intervals)\n",
    "preds = nf.predict(futr_df=AirPassengersPanel_test, level=[90])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a2bd299",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test cross validation can support conformal prediction\n",
    "prediction_intervals = PredictionIntervals()\n",
    "\n",
    "# refit=False, no conformal predictions outputs\n",
    "nf = NeuralForecast(models=[NHITS(h=12, input_size=24, max_steps=1)], freq='M')\n",
    "test_fail(\n",
    "    nf.cross_validation, \n",
    "    \"Passing prediction_intervals and/or level is only supported with refit=True.\",\n",
    "    args=(AirPassengersPanel_train, prediction_intervals, [30, 70]))\n",
    "\n",
    "# refit=True, we have conformal predictions outputs\n",
    "cv2 = nf.cross_validation(\n",
    "    AirPassengersPanel_train, \n",
    "    prediction_intervals=prediction_intervals,\n",
    "    refit=True,\n",
    "    level=[30, 70]\n",
    ")\n",
    "assert all([col in cv2.columns for col in ['NHITS-lo-30', 'NHITS-hi-30']])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b82e7c70",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# Test quantile and level argument in predict for different models and errors\n",
    "prediction_intervals = PredictionIntervals(method=\"conformal_error\")\n",
    "\n",
    "models = []\n",
    "for nf_model in [NHITS, LSTM, TSMixer]:\n",
    "    params = {\"h\": 12, \"input_size\": 24, \"max_steps\": 1, \"loss\": MAE()}\n",
    "    if nf_model.__name__ == \"TSMixer\":\n",
    "        params.update({\"n_series\": 2})\n",
    "    models.append(nf_model(**params))\n",
    "\n",
    "    params = {\"h\": 12, \"input_size\": 24, \"max_steps\": 1, \"loss\": DistributionLoss(distribution=\"Normal\")}\n",
    "    if nf_model.__name__ == \"TSMixer\":\n",
    "        params.update({\"n_series\": 2})\n",
    "    models.append(nf_model(**params))\n",
    "\n",
    "    params = {\"h\": 12, \"input_size\": 24, \"max_steps\": 1, \"loss\": IQLoss()}\n",
    "    if nf_model.__name__ == \"TSMixer\":\n",
    "        params.update({\"n_series\": 2})\n",
    "    models.append(nf_model(**params))\n",
    "\n",
    "nf = NeuralForecast(models=models, freq='M')\n",
    "nf.fit(AirPassengersPanel_train, prediction_intervals=prediction_intervals)\n",
    "# Test default prediction\n",
    "preds = nf.predict(futr_df=AirPassengersPanel_test)\n",
    "assert list(preds.columns) == ['unique_id', 'ds', 'NHITS', 'NHITS1', 'NHITS1-median', 'NHITS1-lo-90',\n",
    "       'NHITS1-lo-80', 'NHITS1-hi-80', 'NHITS1-hi-90', 'NHITS2_ql0.5', 'LSTM',\n",
    "       'LSTM1', 'LSTM1-median', 'LSTM1-lo-90', 'LSTM1-lo-80', 'LSTM1-hi-80',\n",
    "       'LSTM1-hi-90', 'LSTM2_ql0.5', 'TSMixer', 'TSMixer1', 'TSMixer1-median',\n",
    "       'TSMixer1-lo-90', 'TSMixer1-lo-80', 'TSMixer1-hi-80', 'TSMixer1-hi-90',\n",
    "       'TSMixer2_ql0.5']\n",
    "# Test quantile prediction\n",
    "preds = nf.predict(futr_df=AirPassengersPanel_test, quantiles=[0.2, 0.3])\n",
    "assert list(preds.columns) == ['unique_id', 'ds', 'NHITS', 'NHITS-ql0.2', 'NHITS-ql0.3', 'NHITS1',\n",
    "       'NHITS1_ql0.2', 'NHITS1_ql0.3', 'NHITS2_ql0.2', 'NHITS2_ql0.3', 'LSTM',\n",
    "       'LSTM-ql0.2', 'LSTM-ql0.3', 'LSTM1', 'LSTM1_ql0.2', 'LSTM1_ql0.3',\n",
    "       'LSTM2_ql0.2', 'LSTM2_ql0.3', 'TSMixer', 'TSMixer-ql0.2',\n",
    "       'TSMixer-ql0.3', 'TSMixer1', 'TSMixer1_ql0.2', 'TSMixer1_ql0.3',\n",
    "       'TSMixer2_ql0.2', 'TSMixer2_ql0.3']\n",
    "# Test level prediction\n",
    "preds = nf.predict(futr_df=AirPassengersPanel_test, level=[80, 90])\n",
    "assert list(preds.columns) == ['unique_id', 'ds', 'NHITS', 'NHITS-lo-90', 'NHITS-lo-80', 'NHITS-hi-80',\n",
    "       'NHITS-hi-90', 'NHITS1', 'NHITS1-lo-90', 'NHITS1-lo-80', 'NHITS1-hi-80',\n",
    "       'NHITS1-hi-90', 'NHITS2-lo-90', 'NHITS2-lo-80', 'NHITS2-hi-80',\n",
    "       'NHITS2-hi-90', 'LSTM', 'LSTM-lo-90', 'LSTM-lo-80', 'LSTM-hi-80',\n",
    "       'LSTM-hi-90', 'LSTM1', 'LSTM1-lo-90', 'LSTM1-lo-80', 'LSTM1-hi-80',\n",
    "       'LSTM1-hi-90', 'LSTM2-lo-90', 'LSTM2-lo-80', 'LSTM2-hi-80',\n",
    "       'LSTM2-hi-90', 'TSMixer', 'TSMixer-lo-90', 'TSMixer-lo-80',\n",
    "       'TSMixer-hi-80', 'TSMixer-hi-90', 'TSMixer1', 'TSMixer1-lo-90',\n",
    "       'TSMixer1-lo-80', 'TSMixer1-hi-80', 'TSMixer1-hi-90', 'TSMixer2-lo-90',\n",
    "       'TSMixer2-lo-80', 'TSMixer2-hi-80', 'TSMixer2-hi-90']\n",
    "# Re-Test default prediction - note that they are different from the first test (this is expected)\n",
    "preds = nf.predict(futr_df=AirPassengersPanel_test)\n",
    "assert list(preds.columns) == ['unique_id', 'ds', 'NHITS', 'NHITS1', 'NHITS1-median', 'NHITS2_ql0.5',\n",
    "       'LSTM', 'LSTM1', 'LSTM1-median', 'LSTM2_ql0.5', 'TSMixer', 'TSMixer1',\n",
    "       'TSMixer1-median', 'TSMixer2_ql0.5']"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
